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 <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(
      mlir::Operation *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
  );

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

  ConstraintDependencyGraph translate(ConstrainRefRemappings translation) const;

  ConstrainRefSet getConstrainingValues(const ConstrainRef &ref) const;
 
  /*
  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), signalSets(other.signalSets),
        constantSets(other.constantSets), tables() {}
  ConstraintDependencyGraph &operator=(const ConstraintDependencyGraph &other) {
    mod = other.mod;
    structDef = other.structDef;
    signalSets = other.signalSets;
    constantSets = other.constantSets;
  }
  ~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;
 
  // 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;
 
  // Also mutable for caching within otherwise const lookup operations.
  mutable mlir::SymbolTableCollection tables;

  ConstraintDependencyGraph(mlir::ModuleOp m, component::StructDefOp s)
      : mod(m), structDef(s), signalSets() {}

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

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

class ConstraintDependencyGraphStructAnalysis
    : public StructAnalysis<ConstraintDependencyGraph, NoContext> {
public:
  using StructAnalysis::StructAnalysis;
 
  mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager, NoContext &_
  ) override {
    return runAnalysis(solver, moduleAnalysisManager);
  }

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

class ConstraintDependencyGraphModuleAnalysis
    : public ModuleAnalysis<
          ConstraintDependencyGraph, NoContext, ConstraintDependencyGraphStructAnalysis> {
 
public:
  using ModuleAnalysis::ModuleAnalysis;
 
protected:
  void initializeSolver(mlir::DataFlowSolver &solver) override {
    (void)solver.load<ConstrainRefAnalysis>();
  }
 
  NoContext getContext() override { return {}; }
};
 
} // namespace llzk