SparseAnalysis.h

Go to the documentation of this file.

//===- SparseAnalysis.h - Sparse data-flow analysis -----------------------===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Adapted from mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h.
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include <mlir/Analysis/DataFlow/SparseAnalysis.h>
#include <mlir/Analysis/DataFlowFramework.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Interfaces/CallInterfaces.h>
#include <mlir/Interfaces/ControlFlowInterfaces.h>
 
#include <llvm/ADT/SmallPtrSet.h>
 
namespace llzk::dataflow {
 
using AbstractSparseLattice = mlir::dataflow::AbstractSparseLattice;
 
//===----------------------------------------------------------------------===//
// AbstractSparseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//

class AbstractSparseForwardDataFlowAnalysis : public mlir::DataFlowAnalysis {
public:
  mlir::LogicalResult initialize(mlir::Operation *top) override;

  mlir::LogicalResult visit(mlir::ProgramPoint *point) override;
 
protected:
  explicit AbstractSparseForwardDataFlowAnalysis(mlir::DataFlowSolver &solver);

  virtual mlir::LogicalResult visitOperationImpl(
      mlir::Operation *op, mlir::ArrayRef<const AbstractSparseLattice *> operandLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) = 0;

  virtual void visitExternalCallImpl(
      mlir::CallOpInterface call, mlir::ArrayRef<const AbstractSparseLattice *> argumentLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) = 0;

  virtual void visitNonControlFlowArgumentsImpl(
      mlir::Operation *op, const mlir::RegionSuccessor &successor,
      mlir::ArrayRef<AbstractSparseLattice *> argLattices, unsigned firstIndex
  ) = 0;

  virtual AbstractSparseLattice *getLatticeElement(mlir::Value value) = 0;

  const AbstractSparseLattice *getLatticeElementFor(mlir::ProgramPoint *point, mlir::Value value);

  virtual void setToEntryState(AbstractSparseLattice *lattice) = 0;
  void setAllToEntryStates(mlir::ArrayRef<AbstractSparseLattice *> lattices);

  void join(AbstractSparseLattice *lhs, const AbstractSparseLattice &rhs);

  mlir::SymbolTableCollection tables;
 
private:
  mlir::LogicalResult initializeRecursively(mlir::Operation *op);

  mlir::LogicalResult visitOperation(mlir::Operation *op);

  void visitBlock(mlir::Block *block);

  void visitRegionSuccessors(
      mlir::ProgramPoint *point, mlir::RegionBranchOpInterface branch,
      mlir::RegionBranchPoint successor, mlir::ArrayRef<AbstractSparseLattice *> lattices
  );
};
 
//===----------------------------------------------------------------------===//
// SparseForwardDataFlowAnalysis
//===----------------------------------------------------------------------===//

template <typename StateT>
class SparseForwardDataFlowAnalysis : public AbstractSparseForwardDataFlowAnalysis {
  static_assert(
      std::is_base_of<AbstractSparseLattice, StateT>::value,
      "analysis state class expected to subclass AbstractSparseLattice"
  );
 
public:
  explicit SparseForwardDataFlowAnalysis(mlir::DataFlowSolver &s)
      : AbstractSparseForwardDataFlowAnalysis(s) {}

  virtual mlir::LogicalResult visitOperation(
      mlir::Operation *op, mlir::ArrayRef<const StateT *> operands, mlir::ArrayRef<StateT *> results
  ) = 0;

  virtual void visitExternalCall(
      mlir::CallOpInterface call, mlir::ArrayRef<const StateT *> argumentLattices,
      mlir::ArrayRef<StateT *> resultLattices
  ) {
    setAllToEntryStates(resultLattices);
  }

  virtual void visitNonControlFlowArguments(
      mlir::Operation *op, const mlir::RegionSuccessor &successor,
      mlir::ArrayRef<StateT *> argLattices, unsigned firstIndex
  ) {
    setAllToEntryStates(argLattices.take_front(firstIndex));
    setAllToEntryStates(argLattices.drop_front(firstIndex + successor.getSuccessorInputs().size()));
  }
 
protected:
  StateT *getLatticeElement(mlir::Value value) override { return getOrCreate<StateT>(value); }

  const StateT *getLatticeElementFor(mlir::ProgramPoint *point, mlir::Value value) {
    return static_cast<const StateT *>(
        AbstractSparseForwardDataFlowAnalysis::getLatticeElementFor(point, value)
    );
  }

  virtual void setToEntryState(StateT *lattice) = 0;
  void setAllToEntryStates(mlir::ArrayRef<StateT *> lattices) {
    AbstractSparseForwardDataFlowAnalysis::setAllToEntryStates(
        {reinterpret_cast<AbstractSparseLattice *const *>(lattices.begin()), lattices.size()}
    );
  }
 
private:
  llvm::LogicalResult visitOperationImpl(
      mlir::Operation *op, mlir::ArrayRef<const AbstractSparseLattice *> operandLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) override {
    return visitOperation(
        op,
        {reinterpret_cast<const StateT *const *>(operandLattices.begin()), operandLattices.size()},
        {reinterpret_cast<StateT *const *>(resultLattices.begin()), resultLattices.size()}
    );
  }
  void visitExternalCallImpl(
      mlir::CallOpInterface call, mlir::ArrayRef<const AbstractSparseLattice *> argumentLattices,
      mlir::ArrayRef<AbstractSparseLattice *> resultLattices
  ) override {
    visitExternalCall(
        call,
        {reinterpret_cast<const StateT *const *>(argumentLattices.begin()),
         argumentLattices.size()},
        {reinterpret_cast<StateT *const *>(resultLattices.begin()), resultLattices.size()}
    );
  }
  void visitNonControlFlowArgumentsImpl(
      mlir::Operation *op, const mlir::RegionSuccessor &successor,
      mlir::ArrayRef<AbstractSparseLattice *> argLattices, unsigned firstIndex
  ) override {
    visitNonControlFlowArguments(
        op, successor, {reinterpret_cast<StateT *const *>(argLattices.begin()), argLattices.size()},
        firstIndex
    );
  }
  void setToEntryState(AbstractSparseLattice *lattice) override {
    return setToEntryState(reinterpret_cast<StateT *>(lattice));
  }
};
 
} // namespace llzk::dataflow