ConstraintDependencyGraph.h

Go to the documentation of this file.

//===-- ConstraintDependencyGraph.h -----------------------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Analysis/AnalysisWrappers.h"
#include "llzk/Analysis/SourceRef.h"
#include "llzk/Analysis/SourceRefLattice.h"
#include "llzk/Analysis/SparseAnalysis.h"
#include "llzk/Dialect/Array/IR/Ops.h"
 
#include <mlir/Pass/AnalysisManager.h>
 
#include <llvm/ADT/EquivalenceClasses.h>
 
namespace mlir {
 
class DataFlowSolver;
 
} // namespace mlir
 
namespace llzk {
 
using SourceRefRemappings = std::vector<std::pair<SourceRef, SourceRefLatticeValue>>;

class SourceRefAnalysis : public dataflow::SparseForwardDataFlowAnalysis<SourceRefLattice> {
public:
  using Lattice = SourceRefLattice;
  using OperandValues = mlir::DenseMap<mlir::Value, const Lattice *>;
  using Base = dataflow::SparseForwardDataFlowAnalysis<Lattice>;
  using Base::SparseForwardDataFlowAnalysis;
 
  static const Lattice *getLattice(mlir::DataFlowSolver &solver, mlir::Value val);
  static SourceRefLatticeValue getValueState(mlir::DataFlowSolver &solver, mlir::Value val);
  static mlir::FailureOr<SourceRefLatticeValue>
  getWriteTargetState(mlir::DataFlowSolver &solver, mlir::Operation *op);

  mlir::LogicalResult visitOperation(
      mlir::Operation *op, mlir::ArrayRef<const Lattice *> operands,
      mlir::ArrayRef<Lattice *> results
  ) override;
 
  void visitExternalCall(
      mlir::CallOpInterface call, mlir::ArrayRef<const Lattice *> argumentLattices,
      mlir::ArrayRef<Lattice *> resultLattices
  ) override;
 
protected:
  void setToEntryState(Lattice *lattice) override;
 
  // Perform a standard union of operands into the results value.
  static mlir::ChangeResult fallbackOpUpdate(
      mlir::Operation *op, const OperandValues &operandVals, mlir::ArrayRef<Lattice *> results
  );
 
  // Create the references for either an array.read op or an array.extract op, which
  // operate very similarly: index into the first operand using a variable number
  // of provided indices.
  static SourceRefLatticeValue
  arraySubdivisionOpUpdate(array::ArrayAccessOpInterface op, const OperandValues &operandVals);
 
private:
  mlir::SymbolTableCollection tables;
};

struct CDGAnalysisContext {
  bool runIntraprocedural = false;
 
  bool runIntraproceduralAnalysis() const { return runIntraprocedural; }
 
  friend bool operator==(const CDGAnalysisContext &a, const CDGAnalysisContext &b) = default;
};
 
} // namespace llzk
 
template <> struct std::hash<llzk::CDGAnalysisContext> {
  size_t operator()(const llzk::CDGAnalysisContext &c) const {
    return std::hash<bool> {}(c.runIntraprocedural);
  }
};
 
namespace llzk {

class ConstraintDependencyGraph {
public:
  static mlir::FailureOr<ConstraintDependencyGraph> compute(
      mlir::ModuleOp mod, component::StructDefOp s, mlir::DataFlowSolver &solver,
      mlir::AnalysisManager &am, const CDGAnalysisContext &ctx
  );

  void dump() const;
  void print(mlir::raw_ostream &os) const;

  ConstraintDependencyGraph translate(SourceRefRemappings translation) const;

  SourceRefSet getConstrainingValues(const SourceRef &ref) const;
 
  const SourceRefLattice::Ref2Val &getRef2Val() const { return ref2Val; }
 
  /*
  Rule of three, needed for the mlir::SymbolTableCollection, which has no copy constructor.
  Since the mlir::SymbolTableCollection is a caching mechanism, we simply allow default, empty
  construction for copies.
  */
 
  ConstraintDependencyGraph(const ConstraintDependencyGraph &other)
      : mod(other.mod), structDef(other.structDef), ctx(other.ctx), signalSets(other.signalSets),
        constantSets(other.constantSets), ref2Val(other.ref2Val), tables() {}
 
  ConstraintDependencyGraph &operator=(const ConstraintDependencyGraph &other) {
    mod = other.mod;
    structDef = other.structDef;
    ctx = other.ctx;
    signalSets = other.signalSets;
    constantSets = other.constantSets;
    ref2Val = other.ref2Val;
    return *this;
  }
  virtual ~ConstraintDependencyGraph() = default;
 
private:
  mlir::ModuleOp mod;
  // Using mutable because many operations are not const by default, even for "const"-like
  // operations, like "getName()", and this reduces const_casts.
  mutable component::StructDefOp structDef;
  CDGAnalysisContext ctx;
 
  // Transitive closure only over signals.
  llvm::EquivalenceClasses<SourceRef> signalSets;
  // A simple set mapping of constants, as we do not want to compute a transitive closure over
  // constants.
  std::unordered_map<SourceRef, SourceRefSet, SourceRef::Hash> constantSets;
 
  // Maps references to the values where they are found.
  SourceRefLattice::Ref2Val ref2Val;
 
  // Also mutable for caching within otherwise const lookup operations.
  mutable mlir::SymbolTableCollection tables;

  ConstraintDependencyGraph(mlir::ModuleOp m, component::StructDefOp s, const CDGAnalysisContext &c)
      : mod(m), structDef(s), ctx(c) {}

  mlir::LogicalResult computeConstraints(mlir::DataFlowSolver &solver, mlir::AnalysisManager &am);

  void walkConstrainOp(mlir::DataFlowSolver &solver, mlir::Operation *emitOp);
};

class ConstraintDependencyGraphStructAnalysis
    : public StructAnalysis<ConstraintDependencyGraph, CDGAnalysisContext> {
public:
  using StructAnalysis::StructAnalysis;
  ~ConstraintDependencyGraphStructAnalysis() override = default;

  mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager,
      const CDGAnalysisContext &ctx
  ) override;
};

class ConstraintDependencyGraphModuleAnalysis
    : public ModuleAnalysis<
          ConstraintDependencyGraph, CDGAnalysisContext, ConstraintDependencyGraphStructAnalysis> {
 
public:
  // We set the SourceRef analysis as intraprocedural so that calls are treated as "external"
  // calls, and we can leverage the `visitExternalCall` hook to translate `SourceRef`s
  // between function contexts.
  ConstraintDependencyGraphModuleAnalysis(mlir::Operation *op)
      : ModuleAnalysis(op, mlir::DataFlowConfig().setInterprocedural(false)) {}
 
  ~ConstraintDependencyGraphModuleAnalysis() override = default;
 
  void setIntraprocedural(bool runIntraprocedural) {
    ctx = {.runIntraprocedural = runIntraprocedural};
  }
 
protected:
  void initializeSolver() override { (void)solver.load<SourceRefAnalysis>(); }
 
  const CDGAnalysisContext &getContext() const override { return ctx; }
 
private:
  // This "intraprocedural" option is related to the CDG construction, not the dataflow analysis
  // itself.
  CDGAnalysisContext ctx = {.runIntraprocedural = false};
};
 
} // namespace llzk