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/ConstrainRef.h"
#include "llzk/Analysis/ConstrainRefLattice.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 ConstrainRefRemappings = std::vector<std::pair<ConstrainRef, ConstrainRefLatticeValue>>;

class ConstrainRefAnalysis : public dataflow::DenseForwardDataFlowAnalysis<ConstrainRefLattice> {
public:
  using dataflow::DenseForwardDataFlowAnalysis<ConstrainRefLattice>::DenseForwardDataFlowAnalysis;
 
  void visitCallControlFlowTransfer(
      mlir::CallOpInterface call, dataflow::CallControlFlowAction action,
      const ConstrainRefLattice &before, ConstrainRefLattice *after
  ) override;

  void visitOperation(
      mlir::Operation *op, const ConstrainRefLattice &before, ConstrainRefLattice *after
  ) override;
 
protected:
  void setToEntryState(ConstrainRefLattice *lattice) override {
    // the entry state is empty, so do nothing.
  }
 
  // Perform a standard union of operands into the results value.
  mlir::ChangeResult fallbackOpUpdate(
      mlir::Operation *op, const ConstrainRefLattice::ValueMap &operandVals,
      const ConstrainRefLattice &before, ConstrainRefLattice *after
  );
 
  // Perform the update for either a array.read op or an array.extract op, which
  // operate very similarly: index into the first operand using a variable number
  // of provided indices.
  void arraySubdivisionOpUpdate(
      array::ArrayAccessOpInterface op, const ConstrainRefLattice::ValueMap &operandVals,
      const ConstrainRefLattice &before, ConstrainRefLattice *after
  );
 
private:
  mlir::SymbolTableCollection tables;
};

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

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

  ConstraintDependencyGraph translate(ConstrainRefRemappings translation) const;

  ConstrainRefSet getConstrainingValues(const ConstrainRef &ref) const;
 
  const ConstrainRefLattice::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), runIntraprocedural(other.runIntraprocedural),
        signalSets(other.signalSets), constantSets(other.constantSets), ref2Val(other.ref2Val),
        tables() {}
 
  ConstraintDependencyGraph &operator=(const ConstraintDependencyGraph &other) {
    mod = other.mod;
    structDef = other.structDef;
    runIntraprocedural = other.runIntraprocedural;
    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;
  bool runIntraprocedural;
 
  // Transitive closure only over signals.
  llvm::EquivalenceClasses<ConstrainRef> signalSets;
  // A simple set mapping of constants, as we do not want to compute a transitive closure over
  // constants.
  std::unordered_map<ConstrainRef, ConstrainRefSet, ConstrainRef::Hash> constantSets;
 
  // Maps references to the values where they are found.
  ConstrainRefLattice::Ref2Val ref2Val;
 
  // Also mutable for caching within otherwise const lookup operations.
  mutable mlir::SymbolTableCollection tables;

  ConstraintDependencyGraph(
      mlir::ModuleOp m, component::StructDefOp s, bool intraprocedural = false
  )
      : mod(m), structDef(s), runIntraprocedural(intraprocedural) {}

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

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

struct CDGAnalysisContext {
  bool runIntraprocedural = false;
 
  bool runIntraproceduralAnalysis() const { return runIntraprocedural; }
};

class ConstraintDependencyGraphStructAnalysis
    : public StructAnalysis<ConstraintDependencyGraph, CDGAnalysisContext> {
public:
  using StructAnalysis::StructAnalysis;
  virtual ~ConstraintDependencyGraphStructAnalysis() = default;
 
  mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager,
      CDGAnalysisContext &ctx
  ) override {
    return runAnalysis(solver, moduleAnalysisManager, ctx.runIntraproceduralAnalysis());
  }

  mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager,
      bool runIntraprocedural
  );
};

class ConstraintDependencyGraphModuleAnalysis
    : public ModuleAnalysis<
          ConstraintDependencyGraph, CDGAnalysisContext, ConstraintDependencyGraphStructAnalysis> {
 
public:
  using ModuleAnalysis::ModuleAnalysis;
  virtual ~ConstraintDependencyGraphModuleAnalysis() = default;
 
  void setIntraprocedural(bool runIntraprocedural) { intraprocedural = runIntraprocedural; }
 
protected:
  void initializeSolver() override { (void)solver.load<ConstrainRefAnalysis>(); }
 
  CDGAnalysisContext getContext() override { return {.runIntraprocedural = intraprocedural}; }
 
private:
  bool intraprocedural = false;
};
 
} // namespace llzk