diff --git a/changelogs/unreleased/iangneal__analysis-tests.yaml b/changelogs/unreleased/iangneal__analysis-tests.yaml new file mode 100644 index 000000000..8bf1af3f5 --- /dev/null +++ b/changelogs/unreleased/iangneal__analysis-tests.yaml @@ -0,0 +1,6 @@ +added: + - PredecessorAnalysis (-llzk-predecessor-analysis) for testing +changed: + - Implement `CallOpInterface` methods for `CallOp` +removed: + - LLZK port of MLIR DenseAnalysis diff --git a/include/llzk/Analysis/AnalysisPasses.h b/include/llzk/Analysis/AnalysisPasses.h index 0e41ee758..4cfa9519f 100644 --- a/include/llzk/Analysis/AnalysisPasses.h +++ b/include/llzk/Analysis/AnalysisPasses.h @@ -26,6 +26,8 @@ std::unique_ptr createSymbolDefTreePrinterPass(); std::unique_ptr createSymbolUseGraphPrinterPass(); +std::unique_ptr createPredecessorPrinterPass(); + #define GEN_PASS_REGISTRATION #include "llzk/Analysis/AnalysisPasses.h.inc" diff --git a/include/llzk/Analysis/AnalysisPasses.td b/include/llzk/Analysis/AnalysisPasses.td index 5cea03347..0fb0e042d 100644 --- a/include/llzk/Analysis/AnalysisPasses.td +++ b/include/llzk/Analysis/AnalysisPasses.td @@ -108,4 +108,14 @@ def SymbolUseGraphPrinterPass : LLZKPass<"llzk-print-symbol-use-graph"> { let options = [OutputStreamOption, SaveDotGraphOption]; } +def PredecessorPrinterPass : LLZKPass<"llzk-print-predecessors"> { + let summary = "Print the predecessors of all operations."; + let constructor = "llzk::createPredecessorPrinterPass()"; + let options = [OutputStreamOption, + Option<"preRunRequiredAnalyses", "prerun", "bool", + /* default */ "false", + "Whether to pre-run the required dataflow analyses " + "(e.g., liveness analysis).">]; +} + #endif // LLZK_ANALYSIS_TD diff --git a/include/llzk/Analysis/AnalysisUtil.h b/include/llzk/Analysis/AnalysisUtil.h index 7dbc9d244..45ea7f304 100644 --- a/include/llzk/Analysis/AnalysisUtil.h +++ b/include/llzk/Analysis/AnalysisUtil.h @@ -13,13 +13,25 @@ namespace llzk::dataflow { -/// LLZK: Added this utility to ensure analysis is performed for all structs -/// in a given module. -/// -/// @brief Mark all operations from the top and included in the top operation -/// as live so the solver will perform dataflow analyses. +/// @brief Loads analyses required to initialize the Executable and PredecessorState +/// analysis states, which are required for the MLIR Dataflow analyses to properly +/// traverse the LLZK IR. /// @param solver The solver. -/// @param top The top-level operation. -void markAllOpsAsLive(mlir::DataFlowSolver &solver, mlir::Operation *top); +void loadRequiredAnalyses(mlir::DataFlowSolver &solver); + +/// @brief Loads and runs analyses required to initialize the Executable and PredecessorState +/// analysis states, which are required for the MLIR Dataflow analyses to properly +/// traverse the LLZK IR. +/// This function pre-runs the analyses, which is helpful in cases where early +/// region-op-body traversal is desired. +/// - the bodies of scf.for, scf.while, scf.if are usually not marked as live +/// initially, so the dataflow analysis will traverse all ops in a function +/// before traversing the insides of region ops. +/// - by pre-running the analysis, the region bodies will be explored as encountered +/// if they are marked as "live" by the dead code analysis. +/// @param solver The solver. +/// @param op The operation to pre-run the analyses on. +/// @return Whether the pre-run analysis was successful. +mlir::LogicalResult loadAndRunRequiredAnalyses(mlir::DataFlowSolver &solver, mlir::Operation *op); } // namespace llzk::dataflow diff --git a/include/llzk/Analysis/AnalysisWrappers.h b/include/llzk/Analysis/AnalysisWrappers.h index 00d3075d7..7d8aae93e 100644 --- a/include/llzk/Analysis/AnalysisWrappers.h +++ b/include/llzk/Analysis/AnalysisWrappers.h @@ -212,7 +212,8 @@ class ModuleAnalysis { /// in the `ModuleOp` that is being subjected to this analysis. /// @param am The module's analysis manager. void constructChildAnalyses(mlir::AnalysisManager &am) { - dataflow::markAllOpsAsLive(solver, modOp); + auto init = dataflow::loadAndRunRequiredAnalyses(solver, modOp); + ensure(init.succeeded(), "solver failed to run on module!"); // The analysis is run at the module level so that lattices are computed // for global functions as well. diff --git a/include/llzk/Analysis/ConstraintDependencyGraph.h b/include/llzk/Analysis/ConstraintDependencyGraph.h index aee09b934..79b46f244 100644 --- a/include/llzk/Analysis/ConstraintDependencyGraph.h +++ b/include/llzk/Analysis/ConstraintDependencyGraph.h @@ -32,12 +32,13 @@ using SourceRefRemappings = std::vector { +class SourceRefAnalysis : public mlir::dataflow::DenseForwardDataFlowAnalysis { public: - using dataflow::DenseForwardDataFlowAnalysis::DenseForwardDataFlowAnalysis; + using mlir::dataflow::DenseForwardDataFlowAnalysis< + SourceRefLattice>::DenseForwardDataFlowAnalysis; void visitCallControlFlowTransfer( - mlir::CallOpInterface call, dataflow::CallControlFlowAction action, + mlir::CallOpInterface call, mlir::dataflow::CallControlFlowAction action, const SourceRefLattice &before, SourceRefLattice *after ) override; diff --git a/include/llzk/Analysis/DenseAnalysis.h b/include/llzk/Analysis/DenseAnalysis.h deleted file mode 100644 index 1d1bbc0d1..000000000 --- a/include/llzk/Analysis/DenseAnalysis.h +++ /dev/null @@ -1,289 +0,0 @@ -//===-- DenseAnalysis.h - Dense data-flow analysis --------------*- 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 -// -// Adapted from mlir/include/mlir/Analysis/DataFlow/DenseAnalysis.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 -// -//===----------------------------------------------------------------------===// -/// -/// \file -/// This file implements (LLZK-tailored) dense data-flow analysis using the -/// data-flow analysis framework. The analysis is forward and conditional and -/// uses the results of dead code analysis to prune dead code during the -/// analysis. -/// -/// This file has been ported from the MLIR dense analysis so that it may be -/// tailored to work for LLZK modules, -/// as LLZK modules have different symbol lookup mechanisms that are currently -/// incompatible with the builtin MLIR dataflow analyses. -/// This file is mostly left as original in MLIR, with notes added where -/// changes have been made. -//===----------------------------------------------------------------------===// - -#pragma once - -#include -#include -#include -#include -#include - -namespace llzk::dataflow { - -//===----------------------------------------------------------------------===// -// AbstractDenseForwardDataFlowAnalysis -//===----------------------------------------------------------------------===// - -using AbstractDenseLattice = mlir::dataflow::AbstractDenseLattice; -using CallControlFlowAction = mlir::dataflow::CallControlFlowAction; - -/// LLZK: This class has been ported from the MLIR DenseAnalysis utilities to -/// allow for the use of custom LLZK symbol lookup logic. The class has been -/// left as unmodified as possible, with explicit comments added where modifications -/// have been made. -/// -/// Base class for dense forward data-flow analyses. Dense data-flow analysis -/// attaches a lattice to program points and implements a transfer function from -/// the lattice before each operation to the lattice after. The lattice contains -/// information about the state of the program at that program point. -/// -/// Visit a program point in forward dense data-flow analysis will invoke the -/// transfer function of the operation preceding the program point iterator. -/// Visit a program point at the begining of block will visit the block itself. -class AbstractDenseForwardDataFlowAnalysis : public mlir::DataFlowAnalysis { -public: - using mlir::DataFlowAnalysis::DataFlowAnalysis; - - /// Initialize the analysis by visiting every program point whose execution - /// may modify the program state; that is, every operation and block. - mlir::LogicalResult initialize(mlir::Operation *top) override; - - /// Visit a program point that modifies the state of the program. If the - /// program point is at the beginning of a block, then the state is propagated - /// from control-flow predecessors or callsites. If the operation before - /// program point iterator is a call operation or region control-flow - /// operation, then the state after the execution of the operation is set by - /// control-flow or the callgraph. Otherwise, this function invokes the - /// operation transfer function before the program point iterator. - mlir::LogicalResult visit(mlir::ProgramPoint *point) override; - -protected: - /// Propagate the dense lattice before the execution of an operation to the - /// lattice after its execution. - virtual mlir::LogicalResult visitOperationImpl( - mlir::Operation *op, const AbstractDenseLattice &before, AbstractDenseLattice *after - ) = 0; - - /// Get the dense lattice on the given lattice anchor. - virtual AbstractDenseLattice *getLattice(mlir::LatticeAnchor anchor) = 0; - - /// Get the dense lattice on the given lattice anchor and add dependent as its - /// dependency. That is, every time the lattice after anchor is updated, the - /// dependent program point must be visited, and the newly triggered visit - /// might update the lattice on dependent. - const AbstractDenseLattice * - getLatticeFor(mlir::ProgramPoint *dependent, mlir::LatticeAnchor anchor); - - /// Set the dense lattice at control flow entry point and propagate an update - /// if it changed. - virtual void setToEntryState(AbstractDenseLattice *lattice) = 0; - - /// Join a lattice with another and propagate an update if it changed. - void join(AbstractDenseLattice *lhs, const AbstractDenseLattice &rhs) { - propagateIfChanged(lhs, lhs->join(rhs)); - } - - /// Visit an operation. If this is a call operation or region control-flow - /// operation, then the state after the execution of the operation is set by - /// control-flow or the callgraph. Otherwise, this function invokes the - /// operation transfer function. - virtual mlir::LogicalResult processOperation(mlir::Operation *op); - - /// Propagate the dense lattice forward along the control flow edge from - /// `regionFrom` to `regionTo` regions of the `branch` operation. `nullopt` - /// values correspond to control flow branches originating at or targeting the - /// `branch` operation itself. Default implementation just joins the states, - /// meaning that operations implementing `RegionBranchOpInterface` don't have - /// any effect on the lattice that isn't already expressed by the interface - /// itself. - virtual void visitRegionBranchControlFlowTransfer( - mlir::RegionBranchOpInterface /*branch*/, std::optional regionFrom, - std::optional regionTo, const AbstractDenseLattice &before, - AbstractDenseLattice *after - ) { - join(after, before); - } - - /// Propagate the dense lattice forward along the call control flow edge, - /// which can be either entering or exiting the callee. Default implementation - /// for enter and exit callee actions just meets the states, meaning that - /// operations implementing `CallOpInterface` don't have any effect on the - /// lattice that isn't already expressed by the interface itself. Default - /// implementation for the external callee action additionally sets the - /// "after" lattice to the entry state. - virtual void visitCallControlFlowTransfer( - mlir::CallOpInterface /*call*/, CallControlFlowAction action, - const AbstractDenseLattice &before, AbstractDenseLattice *after - ) { - join(after, before); - // Note that `setToEntryState` may be a "partial fixpoint" for some - // lattices, e.g., lattices that are lists of maps of other lattices will - // only set fixpoint for "known" lattices. - if (action == CallControlFlowAction::ExternalCallee) { - setToEntryState(after); - } - } - - /// Visit a program point within a region branch operation with predecessors - /// in it. This can either be an entry block of one of the regions of the - /// parent operation itself. - void visitRegionBranchOperation( - mlir::ProgramPoint *point, mlir::RegionBranchOpInterface branch, AbstractDenseLattice *after - ); - - /// LLZK: Added for use of symbol helper caching. - mlir::SymbolTableCollection tables; - -private: - /// Visit a block. The state at the start of the block is propagated from - /// control-flow predecessors or callsites. - void visitBlock(mlir::Block *block); - - /// Visit an operation for which the data flow is described by the - /// `CallOpInterface`. - void visitCallOperation( - mlir::CallOpInterface call, const AbstractDenseLattice &before, AbstractDenseLattice *after - ); -}; - -//===----------------------------------------------------------------------===// -// DenseForwardDataFlowAnalysis -//===----------------------------------------------------------------------===// - -/// LLZK: This class has been ported so that it can inherit from our port of -/// the AbstractDenseForwardDataFlowAnalysis above. It is otherwise left unchanged. -/// -/// A dense forward data-flow analysis for propagating lattices before and -/// after the execution of every operation across the IR by implementing -/// transfer functions for operations. -/// -/// `LatticeT` is expected to be a subclass of `AbstractDenseLattice`. -template -class DenseForwardDataFlowAnalysis : public AbstractDenseForwardDataFlowAnalysis { - static_assert( - std::is_base_of::value, - "analysis state class expected to subclass AbstractDenseLattice" - ); - -public: - using AbstractDenseForwardDataFlowAnalysis::AbstractDenseForwardDataFlowAnalysis; - - /// Visit an operation with the dense lattice before its execution. This - /// function is expected to set the dense lattice after its execution and - /// trigger change propagation in case of change. - virtual mlir::LogicalResult - visitOperation(mlir::Operation *op, const LatticeT &before, LatticeT *after) = 0; - - /// Hook for customizing the behavior of lattice propagation along the call - /// control flow edges. Two types of (forward) propagation are possible here: - /// - `action == CallControlFlowAction::Enter` indicates that: - /// - `before` is the state before the call operation; - /// - `after` is the state at the beginning of the callee entry block; - /// - `action == CallControlFlowAction::Exit` indicates that: - /// - `before` is the state at the end of a callee exit block; - /// - `after` is the state after the call operation. - /// By default, the `after` state is simply joined with the `before` state. - /// Concrete analyses can override this behavior or delegate to the parent - /// call for the default behavior. Specifically, if the `call` op may affect - /// the lattice prior to entering the callee, the custom behavior can be added - /// for `action == CallControlFlowAction::Enter`. If the `call` op may affect - /// the lattice post exiting the callee, the custom behavior can be added for - /// `action == CallControlFlowAction::Exit`. - virtual void visitCallControlFlowTransfer( - mlir::CallOpInterface call, CallControlFlowAction action, const LatticeT &before, - LatticeT *after - ) { - AbstractDenseForwardDataFlowAnalysis::visitCallControlFlowTransfer(call, action, before, after); - } - - /// Hook for customizing the behavior of lattice propagation along the control - /// flow edges between regions and their parent op. The control flows from - /// `regionFrom` to `regionTo`, both of which may be `nullopt` to indicate the - /// parent op. The lattice is propagated forward along this edge. The lattices - /// are as follows: - /// - `before:` - /// - if `regionFrom` is a region, this is the lattice at the end of the - /// block that exits the region; note that for multi-exit regions, the - /// lattices are equal at the end of all exiting blocks, but they are - /// associated with different program points. - /// - otherwise, this is the lattice before the parent op. - /// - `after`: - /// - if `regionTo` is a region, this is the lattice at the beginning of - /// the entry block of that region; - /// - otherwise, this is the lattice after the parent op. - /// By default, the `after` state is simply joined with the `before` state. - /// Concrete analyses can override this behavior or delegate to the parent - /// call for the default behavior. Specifically, if the `branch` op may affect - /// the lattice before entering any region, the custom behavior can be added - /// for `regionFrom == nullopt`. If the `branch` op may affect the lattice - /// after all terminated, the custom behavior can be added for `regionTo == - /// nullptr`. The behavior can be further refined for specific pairs of "from" - /// and "to" regions. - virtual void visitRegionBranchControlFlowTransfer( - mlir::RegionBranchOpInterface branch, std::optional regionFrom, - std::optional regionTo, const LatticeT &before, LatticeT *after - ) { - AbstractDenseForwardDataFlowAnalysis::visitRegionBranchControlFlowTransfer( - branch, regionFrom, regionTo, before, after - ); - } - -protected: - /// Get the dense lattice on this lattice anchor. - LatticeT *getLattice(mlir::LatticeAnchor anchor) override { - return getOrCreate(anchor); - } - - /// Set the dense lattice at control flow entry point and propagate an update - /// if it changed. - virtual void setToEntryState(LatticeT *lattice) = 0; - void setToEntryState(AbstractDenseLattice *lattice) override { - setToEntryState(static_cast(lattice)); - } - - /// Type-erased wrappers that convert the abstract dense lattice to a derived - /// lattice and invoke the virtual hooks operating on the derived lattice. - mlir::LogicalResult visitOperationImpl( - mlir::Operation *op, const AbstractDenseLattice &before, AbstractDenseLattice *after - ) final { - return visitOperation( - op, static_cast(before), static_cast(after) - ); - } - void visitCallControlFlowTransfer( - mlir::CallOpInterface call, CallControlFlowAction action, const AbstractDenseLattice &before, - AbstractDenseLattice *after - ) final { - visitCallControlFlowTransfer( - call, action, static_cast(before), static_cast(after) - ); - } - void visitRegionBranchControlFlowTransfer( - mlir::RegionBranchOpInterface branch, std::optional regionFrom, - std::optional regionTo, const AbstractDenseLattice &before, - AbstractDenseLattice *after - ) final { - visitRegionBranchControlFlowTransfer( - branch, regionFrom, regionTo, static_cast(before), - static_cast(after) - ); - } -}; - -} // namespace llzk::dataflow diff --git a/include/llzk/Analysis/IntervalAnalysis.h b/include/llzk/Analysis/IntervalAnalysis.h index d89e70551..e175c20f1 100644 --- a/include/llzk/Analysis/IntervalAnalysis.h +++ b/include/llzk/Analysis/IntervalAnalysis.h @@ -12,7 +12,6 @@ #include "llzk/Analysis/AbstractLatticeValue.h" #include "llzk/Analysis/AnalysisWrappers.h" #include "llzk/Analysis/ConstraintDependencyGraph.h" -#include "llzk/Analysis/DenseAnalysis.h" #include "llzk/Analysis/Intervals.h" #include "llzk/Analysis/SparseAnalysis.h" #include "llzk/Dialect/Array/IR/Ops.h" @@ -26,6 +25,7 @@ #include "llzk/Util/Compare.h" #include "llzk/Util/Field.h" +#include #include #include #include diff --git a/include/llzk/Analysis/SourceRefLattice.h b/include/llzk/Analysis/SourceRefLattice.h index 7ed45268e..162ceb208 100644 --- a/include/llzk/Analysis/SourceRefLattice.h +++ b/include/llzk/Analysis/SourceRefLattice.h @@ -10,10 +10,11 @@ #pragma once #include "llzk/Analysis/AbstractLatticeValue.h" -#include "llzk/Analysis/DenseAnalysis.h" #include "llzk/Analysis/SourceRef.h" #include "llzk/Util/ErrorHelper.h" +#include + #include namespace llzk { @@ -86,7 +87,7 @@ class SourceRefLatticeValue }; /// A lattice for use in dense analysis. -class SourceRefLattice : public dataflow::AbstractDenseLattice { +class SourceRefLattice : public mlir::dataflow::AbstractDenseLattice { public: // mlir::Value is used for read-like operations that create references in their results, // mlir::Operation* is used for write-like operations that reference values as their destinations diff --git a/include/llzk/Dialect/Function/IR/Ops.td b/include/llzk/Dialect/Function/IR/Ops.td index d7d298de6..a297ec040 100644 --- a/include/llzk/Dialect/Function/IR/Ops.td +++ b/include/llzk/Dialect/Function/IR/Ops.td @@ -365,15 +365,27 @@ def CallOp : FunctionDialectOp< }]>]; let extraClassDeclaration = [{ + //===------------------------------------------------------------------===// + // CallOpInterface Methods + //===------------------------------------------------------------------===// + + ::mlir::Operation *resolveCallableInTable(::mlir::SymbolTableCollection *symbolTable); + + ::mlir::Operation *resolveCallable(); + + //===------------------------------------------------------------------===// + // Utility Methods + //===------------------------------------------------------------------===// + ::mlir::FunctionType getCalleeType(); - /// Return `true` iff the callee function name is `FUNC_NAME_COMPUTE` (this + /// Return `true` iff the callee function name is `FUNC_NAME_COMPUTE` (this /// does not check if the callee function is located within a StructDefOp). inline bool calleeIsCompute() { return FUNC_NAME_COMPUTE == getCallee().getLeafReference(); } - /// Return `true` iff the callee function can contain witness generation code + /// Return `true` iff the callee function can contain witness generation code /// (this does not check if the callee function is located within a StructDefOp) inline bool calleeContainsWitnessGen() { return FUNC_NAME_COMPUTE == getCallee().getLeafReference() || diff --git a/include/llzk/Util/SymbolLookup.h b/include/llzk/Util/SymbolLookup.h index 2284bf3d0..080c8be9f 100644 --- a/include/llzk/Util/SymbolLookup.h +++ b/include/llzk/Util/SymbolLookup.h @@ -101,6 +101,9 @@ class SymbolLookupResultUntyped { } } + /// True iff the symbol is managed (i.e., loaded via an IncludeOp). + bool isManaged() const { return managedResources != nullptr; } + /// Adds a pointer to the set of resources the result has to manage the lifetime of. void manage(mlir::OwningOpRef &&ptr, mlir::SymbolTableCollection &&tables); @@ -158,6 +161,9 @@ template class SymbolLookupResult { bool operator==(const SymbolLookupResult &rhs) const { return inner == rhs.inner; } + /// Return 'true' if the inner resource is managed (i.e., loaded via an IncludeOp). + bool isManaged() const { return inner.isManaged(); } + private: SymbolLookupResultUntyped inner; diff --git a/lib/Analysis/AnalysisUtil.cpp b/lib/Analysis/AnalysisUtil.cpp index a86d0fa11..8aaab3d44 100644 --- a/lib/Analysis/AnalysisUtil.cpp +++ b/lib/Analysis/AnalysisUtil.cpp @@ -9,6 +9,7 @@ #include "llzk/Analysis/AnalysisUtil.h" +#include #include using namespace mlir; @@ -17,16 +18,14 @@ using Executable = mlir::dataflow::Executable; namespace llzk::dataflow { -void markAllOpsAsLive(DataFlowSolver &solver, Operation *top) { - for (Region ®ion : top->getRegions()) { - for (Block &block : region) { - ProgramPoint *point = solver.getProgramPointBefore(&block); - (void)solver.getOrCreateState(point)->setToLive(); - for (Operation &oper : block) { - markAllOpsAsLive(solver, &oper); - } - } - } +void loadRequiredAnalyses(DataFlowSolver &solver) { + solver.load(); + solver.load(); +} + +LogicalResult loadAndRunRequiredAnalyses(DataFlowSolver &solver, Operation *op) { + loadRequiredAnalyses(solver); + return solver.initializeAndRun(op); } } // namespace llzk::dataflow diff --git a/lib/Analysis/CallGraphPasses.cpp b/lib/Analysis/CallGraphPasses.cpp index b4da80e20..601da99b3 100644 --- a/lib/Analysis/CallGraphPasses.cpp +++ b/lib/Analysis/CallGraphPasses.cpp @@ -13,8 +13,8 @@ //===----------------------------------------------------------------------===// /// /// \file -/// This file implements the ` -llzk-print-call-graph` and -/// ` -llzk-print-call-graph-sccs` passes. +/// This file implements the `-llzk-print-call-graph` and +/// `-llzk-print-call-graph-sccs` passes. /// //===----------------------------------------------------------------------===// diff --git a/lib/Analysis/ConstraintDependencyGraph.cpp b/lib/Analysis/ConstraintDependencyGraph.cpp index c7e005204..8eb4dc0ca 100644 --- a/lib/Analysis/ConstraintDependencyGraph.cpp +++ b/lib/Analysis/ConstraintDependencyGraph.cpp @@ -8,7 +8,6 @@ //===----------------------------------------------------------------------===// #include "llzk/Analysis/ConstraintDependencyGraph.h" -#include "llzk/Analysis/DenseAnalysis.h" #include "llzk/Analysis/SourceRefLattice.h" #include "llzk/Dialect/Array/IR/Ops.h" #include "llzk/Dialect/Constrain/IR/Ops.h" @@ -18,6 +17,7 @@ #include "llzk/Util/TypeHelper.h" #include +#include #include #include @@ -39,7 +39,7 @@ using namespace function; /* SourceRefAnalysis */ void SourceRefAnalysis::visitCallControlFlowTransfer( - mlir::CallOpInterface call, dataflow::CallControlFlowAction action, + CallOpInterface call, mlir::dataflow::CallControlFlowAction action, const SourceRefLattice &before, SourceRefLattice *after ) { LLVM_DEBUG(llvm::dbgs() << "SourceRefAnalysis::visitCallControlFlowTransfer: " << call << '\n'); @@ -61,7 +61,7 @@ void SourceRefAnalysis::visitCallControlFlowTransfer( /// `action == CallControlFlowAction::Enter` indicates that: /// - `before` is the state before the call operation; /// - `after` is the state at the beginning of the callee entry block; - if (action == dataflow::CallControlFlowAction::EnterCallee) { + if (action == mlir::dataflow::CallControlFlowAction::EnterCallee) { // We skip updating the incoming lattice for function calls, // as SourceRefs are relative to the containing function/struct, so we don't need to pollute // the callee with the callers values. @@ -74,7 +74,7 @@ void SourceRefAnalysis::visitCallControlFlowTransfer( /// `action == CallControlFlowAction::Exit` indicates that: /// - `before` is the state at the end of a callee exit block; /// - `after` is the state after the call operation. - else if (action == dataflow::CallControlFlowAction::ExitCallee) { + else if (action == mlir::dataflow::CallControlFlowAction::ExitCallee) { // Get the argument values of the lattice by getting the state as it would // have been for the callsite. const SourceRefLattice *beforeCall = getLattice(getProgramPointBefore(call)); @@ -125,8 +125,8 @@ void SourceRefAnalysis::visitCallControlFlowTransfer( } } -mlir::LogicalResult SourceRefAnalysis::visitOperation( - mlir::Operation *op, const SourceRefLattice &before, SourceRefLattice *after +LogicalResult SourceRefAnalysis::visitOperation( + Operation *op, const SourceRefLattice &before, SourceRefLattice *after ) { LLVM_DEBUG(llvm::dbgs() << "SourceRefAnalysis::visitOperation: " << *op << '\n'); // Collect the references that are made by the operands to `op`. @@ -201,11 +201,11 @@ mlir::LogicalResult SourceRefAnalysis::visitOperation( } // Perform a standard union of operands into the results value. -mlir::ChangeResult SourceRefAnalysis::fallbackOpUpdate( - mlir::Operation *op, const SourceRefLattice::ValueMap &operandVals, - const SourceRefLattice &before, SourceRefLattice *after +ChangeResult SourceRefAnalysis::fallbackOpUpdate( + Operation *op, const SourceRefLattice::ValueMap &operandVals, const SourceRefLattice &before, + SourceRefLattice *after ) { - auto updated = mlir::ChangeResult::NoChange; + auto updated = ChangeResult::NoChange; for (auto res : op->getResults()) { auto cur = before.getOrDefault(res); @@ -278,8 +278,8 @@ void SourceRefAnalysis::arraySubdivisionOpUpdate( /* ConstraintDependencyGraph */ -mlir::FailureOr ConstraintDependencyGraph::compute( - mlir::ModuleOp m, StructDefOp s, mlir::DataFlowSolver &solver, mlir::AnalysisManager &am, +FailureOr ConstraintDependencyGraph::compute( + ModuleOp m, StructDefOp s, DataFlowSolver &solver, AnalysisManager &am, const CDGAnalysisContext &ctx ) { ConstraintDependencyGraph cdg(m, s, ctx); diff --git a/lib/Analysis/ConstraintDependencyGraphPass.cpp b/lib/Analysis/ConstraintDependencyGraphPass.cpp index f4fedeffa..3b4ac4802 100644 --- a/lib/Analysis/ConstraintDependencyGraphPass.cpp +++ b/lib/Analysis/ConstraintDependencyGraphPass.cpp @@ -8,7 +8,7 @@ //===----------------------------------------------------------------------===// /// /// \file -/// This file implements the ` -llzk-print-constraint-dependency-graphs` pass. +/// This file implements the `-llzk-print-constraint-dependency-graphs` pass. /// //===----------------------------------------------------------------------===// diff --git a/lib/Analysis/DenseAnalysis.cpp b/lib/Analysis/DenseAnalysis.cpp deleted file mode 100644 index fc41f741c..000000000 --- a/lib/Analysis/DenseAnalysis.cpp +++ /dev/null @@ -1,279 +0,0 @@ -//===- DenseAnalysis.cpp - Dense data-flow analysis -------------*- 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 -// -// Adapted from mlir/lib/Analysis/DataFlow/DenseAnalysis.cpp -// 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 -// -//===----------------------------------------------------------------------===// - -#include "llzk/Analysis/DenseAnalysis.h" -#include "llzk/Dialect/Function/IR/Ops.h" -#include "llzk/Util/ErrorHelper.h" -#include "llzk/Util/SymbolHelper.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include - -using namespace mlir; -using Executable = mlir::dataflow::Executable; -using CFGEdge = mlir::dataflow::CFGEdge; - -namespace llzk { - -using namespace function; - -namespace dataflow { - -//===----------------------------------------------------------------------===// -// AbstractDenseForwardDataFlowAnalysis -//===----------------------------------------------------------------------===// - -LogicalResult AbstractDenseForwardDataFlowAnalysis::initialize(Operation *top) { - // Visit every operation and block. - if (failed(processOperation(top))) { - return failure(); - } - for (Region ®ion : top->getRegions()) { - for (Block &block : region) { - visitBlock(&block); - for (Operation &op : block) { - if (failed(initialize(&op))) { - return failure(); - } - } - } - } - return success(); -} - -LogicalResult AbstractDenseForwardDataFlowAnalysis::visit(ProgramPoint *point) { - if (!point->isBlockStart()) { - return processOperation(point->getPrevOp()); - } - visitBlock(point->getBlock()); - return success(); -} - -/// LLZK: This function has been modified to use LLZK symbol helpers instead of -/// the built-in resolveCallable method. -void AbstractDenseForwardDataFlowAnalysis::visitCallOperation( - CallOpInterface call, const AbstractDenseLattice &before, AbstractDenseLattice *after -) { - // Allow for customizing the behavior of calls to external symbols, including - // when the analysis is explicitly marked as non-interprocedural. - auto callable = resolveCallable(tables, call); - if (!getSolverConfig().isInterprocedural() || - (succeeded(callable) && !callable->get().getCallableRegion())) { - return visitCallControlFlowTransfer(call, CallControlFlowAction::ExternalCallee, before, after); - } - - /// LLZK: The PredecessorState Analysis state does not work for LLZK's custom calls. - /// We therefore accumulate predecessor operations (return ops) manually. - SmallVector predecessors; - callable->get().walk([&predecessors](ReturnOp ret) mutable { predecessors.push_back(ret); }); - - // If we have no predecessors, we cannot reason about dataflow, since there is - // no return value. - if (predecessors.empty()) { - return setToEntryState(after); - } - - for (Operation *predecessor : predecessors) { - // Get the lattices at callee return: - // - // function.def @callee() { - // ... - // return // predecessor - // // latticeAtCalleeReturn - // } - // function.def @caller() { - // ... - // call @callee - // // latticeAfterCall - // ... - // } - AbstractDenseLattice *latticeAfterCall = after; - const AbstractDenseLattice *latticeAtCalleeReturn = - getLatticeFor(getProgramPointAfter(call.getOperation()), getProgramPointAfter(predecessor)); - visitCallControlFlowTransfer( - call, CallControlFlowAction::ExitCallee, *latticeAtCalleeReturn, latticeAfterCall - ); - } -} - -LogicalResult AbstractDenseForwardDataFlowAnalysis::processOperation(Operation *op) { - ProgramPoint *point = getProgramPointAfter(op); - // If the containing block is not executable, bail out. - if (op->getBlock() != nullptr && - !getOrCreateFor(point, getProgramPointBefore(op->getBlock()))->isLive()) { - return success(); - } - - // Get the dense lattice to update. - AbstractDenseLattice *after = getLattice(point); - - // Get the dense state before the execution of the op. - const AbstractDenseLattice *before = getLatticeFor(point, getProgramPointBefore(op)); - - // If this op implements region control-flow, then control-flow dictates its - // transfer function. - if (auto branch = dyn_cast(op)) { - visitRegionBranchOperation(point, branch, after); - return success(); - } - - // If this is a call operation, then join its lattices across known return - // sites. - if (auto call = dyn_cast(op)) { - visitCallOperation(call, *before, after); - return success(); - } - - // Invoke the operation transfer function. - return visitOperationImpl(op, *before, after); -} - -/// LLZK: Removing use of PredecessorState because it does not work with LLZK's -/// CallOp and FuncDefOp definitions. -void AbstractDenseForwardDataFlowAnalysis::visitBlock(Block *block) { - // If the block is not executable, bail out. - ProgramPoint *point = getProgramPointBefore(block); - if (!getOrCreateFor(point, point)->isLive()) { - return; - } - - // Get the dense lattice to update. - AbstractDenseLattice *after = getLattice(point); - - // The dense lattices of entry blocks are set by region control-flow or the - // callgraph. - if (block->isEntryBlock()) { - // Check if this block is the entry block of a callable region. - auto callable = dyn_cast(block->getParentOp()); - if (callable && callable.getCallableRegion() == block->getParent()) { - if (!getSolverConfig().isInterprocedural()) { - return setToEntryState(after); - } - /// LLZK: Get callsites of the callable as the predecessors. - auto moduleOpRes = getTopRootModule(callable.getOperation()); - ensure(succeeded(moduleOpRes), "could not get root module from callable"); - SmallVector callsites; - moduleOpRes->walk([this, &callable, &callsites](CallOp call) mutable { - auto calledFnRes = resolveCallable(tables, call); - if (succeeded(calledFnRes) && - calledFnRes->get().getCallableRegion() == callable.getCallableRegion()) { - callsites.push_back(call); - } - }); - - for (Operation *callsite : callsites) { - // Get the dense lattice before the callsite. - const AbstractDenseLattice *before = getLatticeFor(point, getProgramPointBefore(callsite)); - - visitCallControlFlowTransfer( - llvm::cast(callsite), CallControlFlowAction::EnterCallee, *before, - after - ); - } - return; - } - - // Check if we can reason about the control-flow. - if (auto branch = dyn_cast(block->getParentOp())) { - return visitRegionBranchOperation(point, branch, after); - } - - // Otherwise, we can't reason about the data-flow. - return setToEntryState(after); - } - - // Join the state with the state after the block's predecessors. - for (Block::pred_iterator it = block->pred_begin(), e = block->pred_end(); it != e; ++it) { - // Skip control edges that aren't executable. - Block *predecessor = *it; - if (!getOrCreateFor(point, getLatticeAnchor(predecessor, block)) - ->isLive()) { - continue; - } - - // Merge in the state from the predecessor's terminator. - join(after, *getLatticeFor(point, getProgramPointAfter(predecessor->getTerminator()))); - } -} - -/// LLZK: Removing use of PredecessorState because it does not work with LLZK's lookup logic. -void AbstractDenseForwardDataFlowAnalysis::visitRegionBranchOperation( - ProgramPoint *point, RegionBranchOpInterface branch, AbstractDenseLattice *after -) { - Operation *op = point->isBlockStart() ? point->getBlock()->getParentOp() : point->getPrevOp(); - if (op) { - const AbstractDenseLattice *before; - // If the predecessor is the parent, get the state before the parent. - if (op == branch) { - before = getLatticeFor(point, getProgramPointBefore(op)); - // Otherwise, get the state after the terminator. - } else { - before = getLatticeFor(point, getProgramPointAfter(op)); - } - - // This function is called in two cases: - // 1. when visiting the block (point = block start); - // 2. when visiting the parent operation (point = iter after parent op). - // In both cases, we are looking for predecessor operations of the point, - // 1. predecessor may be the terminator of another block from another - // region (assuming that the block does belong to another region via an - // assertion) or the parent (when parent can transfer control to this - // region); - // 2. predecessor may be the terminator of a block that exits the - // region (when region transfers control to the parent) or the operation - // before the parent. - // In the latter case, just perform the join as it isn't the control flow - // affected by the region. - std::optional regionFrom = - op == branch ? std::optional() : op->getBlock()->getParent()->getRegionNumber(); - if (point->isBlockStart()) { - unsigned regionTo = point->getBlock()->getParent()->getRegionNumber(); - visitRegionBranchControlFlowTransfer(branch, regionFrom, regionTo, *before, after); - } else { - assert(point->getPrevOp() == branch && "expected to be visiting the branch itself"); - // Only need to call the arc transfer when the predecessor is the region - // or the op itself, not the previous op. - if (op->getParentOp() == branch || op == branch) { - visitRegionBranchControlFlowTransfer( - branch, regionFrom, /*regionTo=*/std::nullopt, *before, after - ); - } else { - join(after, *before); - } - } - } -} - -const AbstractDenseLattice * -AbstractDenseForwardDataFlowAnalysis::getLatticeFor(ProgramPoint *dependent, LatticeAnchor anchor) { - AbstractDenseLattice *state = getLattice(anchor); - addDependency(state, dependent); - return state; -} - -} // namespace dataflow - -} // namespace llzk diff --git a/lib/Analysis/IntervalAnalysis.cpp b/lib/Analysis/IntervalAnalysis.cpp index eb4a332bf..8757f7128 100644 --- a/lib/Analysis/IntervalAnalysis.cpp +++ b/lib/Analysis/IntervalAnalysis.cpp @@ -303,37 +303,15 @@ void ExpressionValue::print(mlir::raw_ostream &os) const { /* IntervalAnalysisLattice */ ChangeResult IntervalAnalysisLattice::join(const AbstractSparseLattice &other) { - const auto *rhs = dynamic_cast(&other); - if (!rhs) { - llvm::report_fatal_error("invalid join lattice type"); - } - ChangeResult res = val.update(rhs->getValue()); - for (auto &v : rhs->constraints) { - if (!constraints.contains(v)) { - constraints.insert(v); - res |= ChangeResult::Change; - } - } - return res; + // The update logic is handled in visitOperation; we don't support a generic + // join operation, as it may override valid intervals. + return ChangeResult::NoChange; } ChangeResult IntervalAnalysisLattice::meet(const AbstractSparseLattice &other) { - const auto *rhs = dynamic_cast(&other); - if (!rhs) { - llvm::report_fatal_error("invalid join lattice type"); - } - // Intersect the intervals - ExpressionValue lhsExpr = val.getScalarValue(); - ExpressionValue rhsExpr = rhs->getValue().getScalarValue(); - Interval newInterval = lhsExpr.getInterval().intersect(rhsExpr.getInterval()); - ChangeResult res = setValue(lhsExpr.withInterval(newInterval)); - for (auto &v : rhs->constraints) { - if (!constraints.contains(v)) { - constraints.insert(v); - res |= ChangeResult::Change; - } - } - return res; + // The update logic is handled in visitOperation; we don't support a generic + // meet operation, as it may override valid intervals. + return ChangeResult::NoChange; } void IntervalAnalysisLattice::print(mlir::raw_ostream &os) const { diff --git a/lib/Analysis/PredecessorAnalysisPass.cpp b/lib/Analysis/PredecessorAnalysisPass.cpp new file mode 100644 index 000000000..bcc61c9fd --- /dev/null +++ b/lib/Analysis/PredecessorAnalysisPass.cpp @@ -0,0 +1,247 @@ +//===-- PredecessorPrinterPass.cpp ------------------------------*- C++ -*-===// +// +// Part of the LLZK Project, under the Apache License v2.0. +// See LICENSE.txt for license information. +// Copyright 2026 Project LLZK +// SPDX-License-Identifier: Apache-2.0 +// +//===----------------------------------------------------------------------===// +/// +/// \file +/// This file implements the `-llzk-print-predecessors` pass. +/// +//===----------------------------------------------------------------------===// + +#include "llzk/Analysis/AnalysisPasses.h" +#include "llzk/Analysis/AnalysisUtil.h" +#include "llzk/Dialect/Function/IR/Ops.h" + +#include +#include +#include +#include + +#include +#include +#include + +using namespace mlir; + +namespace llzk { + +using namespace function; + +#define GEN_PASS_DECL_PREDECESSORPRINTERPASS +#define GEN_PASS_DEF_PREDECESSORPRINTERPASS +#include "llzk/Analysis/AnalysisPasses.h.inc" + +/// Prints op without region. +raw_ostream &printRegionless(raw_ostream &os, Operation *op, bool withParent = false) { + std::string s; + llvm::raw_string_ostream ss(s); + if (withParent) { + if (auto fnOp = op->getParentOfType()) { + os << '<' << fnOp.getFullyQualifiedName() << ">:"; + } else { + os << "<(no parent function op)>:"; + } + } + op->print(ss, mlir::OpPrintingFlags().skipRegions()); + ss.flush(); + // Skipping regions inserts a new line we don't want, so trim it here. + llvm::StringRef r(s); + os << r.rtrim(); + return os; +} + +class PredecessorLattice : public mlir::dataflow::AbstractDenseLattice { + /// Maps op -> [predecessor program points] + llvm::MapVector> predecessors; + +public: + using AbstractDenseLattice::AbstractDenseLattice; + + ChangeResult visit(Operation *op, Operation *pred) { + bool newlyInserted = predecessors[op].insert(pred); + return newlyInserted ? ChangeResult::Change : ChangeResult::NoChange; + } + + ChangeResult join(const AbstractDenseLattice &rhs) override { + const auto *other = dynamic_cast(&rhs); + if (!other) { + llvm::report_fatal_error("wrong lattice type provided for join"); + } + ChangeResult r = ChangeResult::NoChange; + for (auto &[op, preds] : other->predecessors) { + for (auto pred : preds) { + r |= visit(op, pred); + } + } + return r; + } + + ChangeResult meet(const AbstractDenseLattice & /*rhs*/) override { + llvm::report_fatal_error("meet operation is not supported for PredecessorLattice"); + return ChangeResult::NoChange; + } + + void print(raw_ostream &os) const override { + if (predecessors.empty()) { + os << "(empty)\n"; + return; + } + for (auto &[k, v] : predecessors) { + os.indent(2); + printRegionless(os, k, true) << " predecessors:"; + llvm::interleave(v, [&os](Operation *p) { + os << '\n'; + printRegionless(os.indent(6), p, true); + }, []() {}); + os << '\n'; + } + } +}; + +class PredecessorAnalysis + : public mlir::dataflow::DenseForwardDataFlowAnalysis { + /// Stream for debug printing. Currently unused, but propagated from `PredecessorPrinterPass` + /// in case it is needed in the future. + [[maybe_unused]] + raw_ostream &os; + + ProgramPoint *getPoint(const PredecessorLattice &l) const { + return dyn_cast(l.getAnchor()); + } + + ChangeResult + updateLattice(Operation *op, const PredecessorLattice &before, PredecessorLattice *after) { + ChangeResult result = after->join(before); + ProgramPoint *pointBefore = getProgramPointBefore(op); + auto *predState = getOrCreate(pointBefore); + if (!predState->getKnownPredecessors().empty()) { + for (Operation *pred : predState->getKnownPredecessors()) { + result |= after->visit(op, pred); + } + } else { + // Predecessor is just the prior or parent op + Operation *pred = pointBefore->isBlockStart() ? op->getParentOp() : pointBefore->getPrevOp(); + result |= after->visit(op, pred); + } + return result; + } + +public: + using Base = DenseForwardDataFlowAnalysis; + + PredecessorAnalysis(DataFlowSolver &s, raw_ostream &ro) : Base(s), os(ro) {} + + LogicalResult visitOperation( + Operation *op, const PredecessorLattice &before, PredecessorLattice *after + ) override { + ChangeResult result = updateLattice(op, before, after); + propagateIfChanged(after, result); + return success(); + } + + void visitCallControlFlowTransfer( + CallOpInterface call, mlir::dataflow::CallControlFlowAction action, + const PredecessorLattice &before, PredecessorLattice *after + ) override { + /// `action == CallControlFlowAction::Enter` indicates that: + /// - `before` is the state before the call operation; + /// - `after` is the state at the beginning of the callee entry block; + if (action == mlir::dataflow::CallControlFlowAction::EnterCallee) { + // We skip updating the incoming lattice for function calls to avoid a + // non-convergence scenario, as calling a function from other contexts + // can cause the lattice values to oscillate and constantly change. + setToEntryState(after); + } + /// `action == CallControlFlowAction::Exit` indicates that: + /// - `before` is the state at the end of a callee exit block; + /// - `after` is the state after the call operation. + else if (action == mlir::dataflow::CallControlFlowAction::ExitCallee) { + // Get the argument values of the lattice by getting the state as it would + // have been for the callsite. + const PredecessorLattice *beforeCall = getLattice(getProgramPointBefore(call)); + ensure(beforeCall, "could not get prior lattice"); + ChangeResult r = after->join(before); + // Perform a visit so that we see the call op in our lattice + r |= updateLattice(call, *beforeCall, after); + propagateIfChanged(after, r); + } + /// `action == CallControlFlowAction::External` indicates that: + /// - `before` is the state before the call operation. + /// - `after` is the state after the call operation, since there is no callee + /// body to enter into. + else if (action == mlir::dataflow::CallControlFlowAction::ExternalCallee) { + // For external calls, we propagate what information we already have from + // before the call to after the call, since the external call won't invalidate + // any of that information. It also, conservatively, makes no assumptions about + // external calls and their computation, so CDG edges will not be computed over + // input arguments to external functions. + join(after, before); + } + } + + void visitRegionBranchControlFlowTransfer( + RegionBranchOpInterface branch, std::optional _regionFrom, + std::optional _regionTo, const PredecessorLattice &before, PredecessorLattice *after + ) override { + // The default implementation is `join(after, before)`, but we want to + // show the predecessor logic for branch operations as well. + (void)visitOperation(branch, before, after); + } + +protected: + void setToEntryState(PredecessorLattice *lattice) override {} +}; + +class PredecessorPrinterPass : public impl::PredecessorPrinterPassBase { + +public: + PredecessorPrinterPass() : impl::PredecessorPrinterPassBase() {} + +protected: + void runOnOperation() override { + markAllAnalysesPreserved(); + // Note: options like `outputStream` are safe to read here, but not in the + // pass constructor. + raw_ostream &os = toStream(outputStream); + + DataFlowSolver solver; + if (preRunRequiredAnalyses) { + ensure( + dataflow::loadAndRunRequiredAnalyses(solver, getOperation()).succeeded(), + "failed to pre-run!" + ); + } else { + dataflow::loadRequiredAnalyses(solver); + } + solver.load(os); + LogicalResult res = solver.initializeAndRun(getOperation()); + + if (res.failed()) { + llvm::report_fatal_error("PredecessorAnalysis failed."); + } + + getOperation()->walk([&](FuncDefOp fnOp) { + Region &fnBody = fnOp.getFunctionBody(); + if (fnBody.empty()) { + return WalkResult::skip(); + } + + ProgramPoint *point = solver.getProgramPointAfter(fnBody.back().getTerminator()); + PredecessorLattice *finalLattice = solver.getOrCreateState(point); + + printRegionless(os, fnOp.getOperation()) << ":\n" << *finalLattice << '\n'; + + return WalkResult::skip(); + }); + } +}; + +std::unique_ptr createPredecessorPrinterPass() { + return std::make_unique(); +} + +} // namespace llzk diff --git a/lib/Analysis/SourceRefLattice.cpp b/lib/Analysis/SourceRefLattice.cpp index 5b6bf8bbc..30433d714 100644 --- a/lib/Analysis/SourceRefLattice.cpp +++ b/lib/Analysis/SourceRefLattice.cpp @@ -8,7 +8,6 @@ //===----------------------------------------------------------------------===// #include "llzk/Analysis/ConstraintDependencyGraph.h" -#include "llzk/Analysis/DenseAnalysis.h" #include "llzk/Analysis/SourceRefLattice.h" #include "llzk/Dialect/Felt/IR/Ops.h" #include "llzk/Dialect/Function/IR/Ops.h" diff --git a/lib/Analysis/SparseAnalysis.cpp b/lib/Analysis/SparseAnalysis.cpp index 7120d684d..fbcc7d421 100644 --- a/lib/Analysis/SparseAnalysis.cpp +++ b/lib/Analysis/SparseAnalysis.cpp @@ -136,47 +136,37 @@ LogicalResult AbstractSparseForwardDataFlowAnalysis::visitOperation(Operation *o operandLattices.push_back(operandLattice); } - // LLZK TODO: Enable for interprocedural analysis. - /* if (auto call = dyn_cast(op)) { - /// LLZK: Use LLZK resolveCallable interface. // If the call operation is to an external function, attempt to infer the // results from the call arguments. - auto callable = resolveCallable(tables, call); - if (!getSolverConfig().isInterprocedural() || - (succeeded(callable) && !callable->get().getCallableRegion())) { + auto callable = dyn_cast_if_present(call.resolveCallable()); + if (!getSolverConfig().isInterprocedural() || (callable && !callable.getCallableRegion())) { visitExternalCallImpl(call, operandLattices, resultLattices); return success(); } // Otherwise, the results of a call operation are determined by the // callgraph. - /// LLZK: The PredecessorState Analysis state does not work for LLZK's custom calls. - /// We therefore accumulate predecessor operations (return ops) manually. - SmallVector predecessors; - callable->get().walk([&predecessors](ReturnOp ret) mutable { predecessors.push_back(ret); }); - + const auto *predecessors = + getOrCreateFor(getProgramPointAfter(op), getProgramPointAfter(call)); // If not all return sites are known, then conservatively assume we can't // reason about the data-flow. - if (predecessors.empty()) { + if (!predecessors->allPredecessorsKnown()) { setAllToEntryStates(resultLattices); return success(); } - for (Operation *predecessor : predecessors) { + for (Operation *predecessor : predecessors->getKnownPredecessors()) { for (auto &&[operand, resLattice] : llvm::zip(predecessor->getOperands(), resultLattices)) { join(resLattice, *getLatticeElementFor(getProgramPointAfter(op), operand)); } } return success(); } - */ // Invoke the operation transfer function. return visitOperationImpl(op, operandLattices, resultLattices); } -/// LLZK: Removing use of PredecessorState because it does not work with LLZK's -/// CallOp and FuncDefOp definitions. void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { // Exit early on blocks with no arguments. if (block->getNumArguments() == 0) { @@ -200,27 +190,17 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { // callgraph. if (block->isEntryBlock()) { // Check if this block is the entry block of a callable region. - // LLZK TODO: Enable for interprocedural analysis. - /* auto callable = dyn_cast(block->getParentOp()); if (callable && callable.getCallableRegion() == block->getParent()) { - /// LLZK: Get callsites of the callable as the predecessors. - auto moduleOpRes = getTopRootModule(callable.getOperation()); - ensure(succeeded(moduleOpRes), "could not get root module from callable"); - SmallVector callsites; - moduleOpRes->walk([this, &callable, &callsites](CallOp call) mutable { - auto calledFnRes = resolveCallable(tables, call); - if (succeeded(calledFnRes) && - calledFnRes->get().getCallableRegion() == callable.getCallableRegion()) { - callsites.push_back(call); - } - }); + const auto *callsites = getOrCreateFor( + getProgramPointBefore(block), getProgramPointAfter(callable) + ); // If not all callsites are known, conservatively mark all lattices as // having reached their pessimistic fixpoints. - if (callsites.empty() || !getSolverConfig().isInterprocedural()) { + if (!callsites->allPredecessorsKnown() || !getSolverConfig().isInterprocedural()) { return setAllToEntryStates(argLattices); } - for (Operation *callsite : callsites) { + for (Operation *callsite : callsites->getKnownPredecessors()) { auto call = cast(callsite); for (auto it : llvm::zip(call.getArgOperands(), argLattices)) { join( @@ -230,7 +210,6 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { } return; } - */ // Check if the lattices can be determined from region control flow. if (auto branch = dyn_cast(block->getParentOp())) { @@ -275,14 +254,14 @@ void AbstractSparseForwardDataFlowAnalysis::visitBlock(Block *block) { } } -/// LLZK: Removing use of PredecessorState because it does not work with LLZK's lookup logic. void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors( ProgramPoint *point, RegionBranchOpInterface branch, RegionBranchPoint successor, ArrayRef lattices ) { - Operation *op = point->isBlockStart() ? point->getBlock()->getParentOp() : point->getPrevOp(); + const auto *predecessors = getOrCreateFor(point, point); + assert(predecessors->allPredecessorsKnown() && "unexpected unresolved region successors"); - if (op) { + for (Operation *op : predecessors->getKnownPredecessors()) { // Get the incoming successor operands. std::optional operands; @@ -299,19 +278,11 @@ void AbstractSparseForwardDataFlowAnalysis::visitRegionSuccessors( return setAllToEntryStates(lattices); } - ValueRange inputs; - - /// LLZK: We only handle these kinds of region ops with inputs for now. - if (auto forOp = dyn_cast(op)) { - inputs = forOp.getRegionIterArgs(); - } else if (auto whileOp = dyn_cast(op)) { - inputs = whileOp.getRegionIterArgs(); - } - - if (inputs.size() != operands->size()) { - // We can't reason about the data-flow. - return setAllToEntryStates(lattices); - } + ValueRange inputs = predecessors->getSuccessorInputs(op); + assert( + inputs.size() == operands->size() && + "expected the same number of successor inputs as operands" + ); unsigned firstIndex = 0; if (inputs.size() != lattices.size()) { diff --git a/lib/Dialect/Function/IR/Ops.cpp b/lib/Dialect/Function/IR/Ops.cpp index d40f091e2..f89c41166 100644 --- a/lib/Dialect/Function/IR/Ops.cpp +++ b/lib/Dialect/Function/IR/Ops.cpp @@ -819,4 +819,19 @@ SmallVector CallOp::toVectorOfValueRange(OperandRangeRange input) { return output; } +Operation *CallOp::resolveCallableInTable(SymbolTableCollection *symbolTable) { + FailureOr> res = + llzk::resolveCallable(*symbolTable, *this); + if (LogicalResult(res).failed() || res->isManaged()) { + // Cannot return pointer to a managed Operation since it would cause memory errors. + return nullptr; + } + return res->get(); +} + +Operation *CallOp::resolveCallable() { + SymbolTableCollection tables; + return resolveCallableInTable(&tables); +} + } // namespace llzk::function diff --git a/test/Analysis/interval_analysis/interval_analysis_pass_compute.llzk b/test/Analysis/interval_analysis/interval_analysis_pass_compute.llzk index 1b421899b..edee5083d 100644 --- a/test/Analysis/interval_analysis/interval_analysis_pass_compute.llzk +++ b/test/Analysis/interval_analysis/interval_analysis_pass_compute.llzk @@ -546,6 +546,8 @@ module attributes {llzk.lang} { %3 = felt.const 3 scf.yield %3 : !felt.type } + // COM: This is currently unsound, as it doesn't account for whether + // or not the loop will run. struct.writem %self[@a] = %c : !struct.type<@YieldForConst>, !felt.type function.return %self : !struct.type<@YieldForConst> } @@ -564,6 +566,39 @@ module attributes {llzk.lang} { // ----- +module attributes {llzk.lang} { + struct.def @YieldForConst { + struct.member @a: !felt.type + + function.def @compute(%x: index) -> !struct.type<@YieldForConst> { + %self = struct.new : !struct.type<@YieldForConst> + %0 = arith.constant 0 : index + %1 = arith.constant 1 : index + %f0 = felt.const 0 + scf.for %i = %0 to %x step %1 { + %3 = felt.const 3 + // COM: This is currently unsound, as it doesn't account for whether + // or not the loop will run. + struct.writem %self[@a] = %3 : !struct.type<@YieldForConst>, !felt.type + scf.yield + } + function.return %self : !struct.type<@YieldForConst> + } + + function.def @constrain(%self: !struct.type<@YieldForConst>, %x: index) { + function.return + } + } +} + +// CHECK-LABEL: @YieldForConst StructIntervals { +// CHECK-NEXT: compute { +// CHECK-NEXT: %arg0 in Entire +// CHECK-NEXT: %self[@a] in Degenerate(3) +// CHECK-NEXT: } + +// ----- + // No guesses will be made about the computed value over the loop module attributes {llzk.lang} { struct.def @YieldForSum { diff --git a/test/Analysis/predecessor_analysis_pass.llzk b/test/Analysis/predecessor_analysis_pass.llzk new file mode 100644 index 000000000..00cc4a8a1 --- /dev/null +++ b/test/Analysis/predecessor_analysis_pass.llzk @@ -0,0 +1,158 @@ +// RUN: llzk-opt -split-input-file --pass-pipeline='builtin.module(llzk-print-predecessors{stream=outs})' %s -o /dev/null | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + function.def @free() { + %c9 = felt.const 7 + function.return + } + + struct.def @A<[@Start, @Stop, @Step]> { + struct.member @foo : !array.type<4 x !felt.type> + function.def @compute(%val: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> { + %self = struct.new : <@A<[@Start, @Stop, @Step]>> + %array = array.new : !array.type<4 x !felt.type> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + %start = poly.read_const @Start : index + %stop = poly.read_const @Stop : index + %step = poly.read_const @Step : index + + scf.for %i = %c0 to %c4 step %c1 { + array.write %array[%c0] = %val : !array.type<4 x !felt.type>, !felt.type + scf.yield + } + + function.call @free() : () -> () + + scf.for %i = %start to %stop step %step { + %v = array.read %array[%c0] : !array.type<4 x !felt.type>, !felt.type + array.write %array[%c1] = %v : !array.type<4 x !felt.type>, !felt.type + scf.yield + } + + struct.writem %self[@foo] = %array : !struct.type<@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> + + function.return %self : !struct.type<@A<[@Start, @Stop, @Step]>> + } + + function.def @constrain(%self: !struct.type<@A<[@Start, @Stop, @Step]>>, %val: !felt.type) { + function.return + } + } +} + +// CHECK-LABEL: function.def @free() {...}: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @free() {...} +// CHECK-NEXT: <@free>:function.return predecessors: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 +// CHECK-LABEL: function.def @compute(%arg0: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> attributes {function.allow_witness} {...}: +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@foo] = %array : <@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> predecessors: +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:function.return %self : !struct.type<@A<[@Start, @Stop, @Step]>> predecessors: +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@foo] = %array : <@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @free() {...} +// CHECK-NEXT: <@free>:function.return predecessors: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 +// CHECK-NEXT: <@A::@compute>:function.call @free() : () -> () predecessors: +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} +// CHECK-NEXT: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} predecessors: +// CHECK-NEXT: <@free>:function.return +// CHECK-NEXT: <@A::@compute>:%3 = array.read %array[%c0] : <4 x !felt.type>, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} +// CHECK-NEXT: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:array.write %array[%c1] = %3 : <4 x !felt.type>, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%3 = array.read %array[%c0] : <4 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:scf.yield predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c1] = %3 : <4 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A<[@Start, @Stop, @Step]>> predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @compute(%arg0: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> attributes {function.allow_witness} {...} +// CHECK-NEXT: <@A::@compute>:%array = array.new : <4 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A<[@Start, @Stop, @Step]>> +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%array = array.new : <4 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%c1 = arith.constant 1 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index +// CHECK-NEXT: <@A::@compute>:%c4 = arith.constant 4 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c1 = arith.constant 1 : index +// CHECK-NEXT: <@A::@compute>:%0 = poly.read_const @Start : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c4 = arith.constant 4 : index +// CHECK-NEXT: <@A::@compute>:%1 = poly.read_const @Stop : index predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = poly.read_const @Start : index +// CHECK-NEXT: <@A::@compute>:%2 = poly.read_const @Step : index predecessors: +// CHECK-NEXT: <@A::@compute>:%1 = poly.read_const @Stop : index +// CHECK-NEXT: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = poly.read_const @Step : index +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %arg0 : <4 x !felt.type>, !felt.type predecessors: +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} +// CHECK-NEXT: <@A::@compute>:scf.yield predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %arg0 : <4 x !felt.type>, !felt.type +// CHECK-LABEL: function.def @constrain(%arg0: !struct.type<@A<[@Start, @Stop, @Step]>>, %arg1: !felt.type) attributes {function.allow_constraint} {...}: +// CHECK-NEXT: <@A::@constrain>:function.return predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @constrain(%arg0: !struct.type<@A<[@Start, @Stop, @Step]>>, %arg1: !felt.type) attributes {function.allow_constraint} {...} + +// ----- + +module attributes {llzk.lang = "llzk"} { + struct.def @A { + struct.member @sums : !array.type<1 x !felt.type> + function.def @compute(%a: !felt.type, %b: !felt.type, %c: i1) -> !struct.type<@A> { + %c0 = arith.constant 0 : index + %self = struct.new : !struct.type<@A> + %sums = array.new : !array.type<1 x !felt.type> + + %ans = scf.if %c -> !felt.type { + %sum = felt.add %a, %b + scf.yield %sum : !felt.type + } else { + %diff = felt.add %a, %a + scf.yield %diff : !felt.type + } + %to_write = felt.sub %ans, %b + array.write %sums[%c0] = %to_write : !array.type<1 x !felt.type>, !felt.type + + struct.writem %self[@sums] = %sums : !struct.type<@A>, !array.type<1 x !felt.type> + function.return %self : !struct.type<@A> + } + + function.def @constrain(%self: !struct.type<@A>, %a: !felt.type, %b: !felt.type, %c: i1) { + function.return + } + } +} + +// CHECK-LABEL: function.def @compute(%arg0: !felt.type, %arg1: !felt.type, %arg2: i1) -> !struct.type<@A> attributes {function.allow_witness} {...}: +// CHECK-NEXT: <@A::@compute>:%1 = felt.sub %0, %arg1 : !felt.type, !felt.type predecessors: +// CHECK-DAG: <@A::@compute>:scf.yield %2 : !felt.type +// CHECK-DAG: <@A::@compute>:scf.yield %2 : !felt.type +// CHECK-DAG: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %1 : <1 x !felt.type>, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%1 = felt.sub %0, %arg1 : !felt.type, !felt.type +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@sums] = %array : <@A>, !array.type<1 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %1 : <1 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:function.return %self : !struct.type<@A> predecessors: +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@sums] = %array : <@A>, !array.type<1 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @compute(%arg0: !felt.type, %arg1: !felt.type, %arg2: i1) -> !struct.type<@A> attributes {function.allow_witness} {...} +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A> predecessors: +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index +// CHECK-NEXT: <@A::@compute>:%array = array.new : <1 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A> +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} predecessors: +// CHECK-NEXT: <@A::@compute>:%array = array.new : <1 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg1 : !felt.type, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} +// CHECK-NEXT: <@A::@compute>:scf.yield %2 : !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg1 : !felt.type, !felt.type +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg0 : !felt.type, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} +// CHECK-NEXT: <@A::@compute>:scf.yield %2 : !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg0 : !felt.type, !felt.type +// CHECK-LABEL: function.def @constrain(%arg0: !struct.type<@A>, %arg1: !felt.type, %arg2: !felt.type, %arg3: i1) attributes {function.allow_constraint} {...}: +// CHECK-NEXT: <@A::@constrain>:function.return predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @constrain(%arg0: !struct.type<@A>, %arg1: !felt.type, %arg2: !felt.type, %arg3: i1) attributes {function.allow_constraint} {...} diff --git a/test/Analysis/predecessor_analysis_prerun_pass.llzk b/test/Analysis/predecessor_analysis_prerun_pass.llzk new file mode 100644 index 000000000..41da46e49 --- /dev/null +++ b/test/Analysis/predecessor_analysis_prerun_pass.llzk @@ -0,0 +1,157 @@ +// RUN: llzk-opt -split-input-file --pass-pipeline='builtin.module(llzk-print-predecessors{stream=outs prerun})' %s -o /dev/null | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + function.def @free() { + %c9 = felt.const 7 + function.return + } + + struct.def @A<[@Start, @Stop, @Step]> { + struct.member @foo : !array.type<4 x !felt.type> + function.def @compute(%val: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> { + %self = struct.new : <@A<[@Start, @Stop, @Step]>> + %array = array.new : !array.type<4 x !felt.type> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + %start = poly.read_const @Start : index + %stop = poly.read_const @Stop : index + %step = poly.read_const @Step : index + + scf.for %i = %c0 to %c4 step %c1 { + array.write %array[%c0] = %val : !array.type<4 x !felt.type>, !felt.type + scf.yield + } + + function.call @free() : () -> () + + scf.for %i = %start to %stop step %step { + %v = array.read %array[%c0] : !array.type<4 x !felt.type>, !felt.type + array.write %array[%c1] = %v : !array.type<4 x !felt.type>, !felt.type + scf.yield + } + + struct.writem %self[@foo] = %array : !struct.type<@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> + + function.return %self : !struct.type<@A<[@Start, @Stop, @Step]>> + } + + function.def @constrain(%self: !struct.type<@A<[@Start, @Stop, @Step]>>, %val: !felt.type) { + function.return + } + } +} + +// CHECK-LABEL: function.def @free() {...}: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @free() {...} +// CHECK-NEXT: <@free>:function.return predecessors: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 +// CHECK-LABEL: function.def @compute(%arg0: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> attributes {function.allow_witness} {...}: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @free() {...} +// CHECK-NEXT: <@free>:function.return predecessors: +// CHECK-NEXT: <@free>:%felt_const_7 = felt.const 7 +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A<[@Start, @Stop, @Step]>> predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @compute(%arg0: !felt.type) -> !struct.type<@A<[@Start, @Stop, @Step]>> attributes {function.allow_witness} {...} +// CHECK-NEXT: <@A::@compute>:%array = array.new : <4 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A<[@Start, @Stop, @Step]>> +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%array = array.new : <4 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%c1 = arith.constant 1 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index +// CHECK-NEXT: <@A::@compute>:%c4 = arith.constant 4 : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c1 = arith.constant 1 : index +// CHECK-NEXT: <@A::@compute>:%0 = poly.read_const @Start : index predecessors: +// CHECK-NEXT: <@A::@compute>:%c4 = arith.constant 4 : index +// CHECK-NEXT: <@A::@compute>:%1 = poly.read_const @Stop : index predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = poly.read_const @Start : index +// CHECK-NEXT: <@A::@compute>:%2 = poly.read_const @Step : index predecessors: +// CHECK-NEXT: <@A::@compute>:%1 = poly.read_const @Stop : index +// CHECK-NEXT: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = poly.read_const @Step : index +// CHECK-NEXT: <@A::@compute>:function.call @free() : () -> () predecessors: +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} predecessors: +// CHECK-NEXT: <@free>:function.return +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@foo] = %array : <@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> predecessors: +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:function.return %self : !struct.type<@A<[@Start, @Stop, @Step]>> predecessors: +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@foo] = %array : <@A<[@Start, @Stop, @Step]>>, !array.type<4 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%3 = array.read %array[%c0] : <4 x !felt.type>, !felt.type predecessors: +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %0 to %1 step %2 {...} +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:array.write %array[%c1] = %3 : <4 x !felt.type>, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%3 = array.read %array[%c0] : <4 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:scf.yield predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c1] = %3 : <4 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %arg0 : <4 x !felt.type>, !felt.type predecessors: +// CHECK-DAG: <@A::@compute>:scf.for %arg1 = %c0 to %c4 step %c1 {...} +// CHECK-DAG: <@A::@compute>:scf.yield +// CHECK-NEXT: <@A::@compute>:scf.yield predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %arg0 : <4 x !felt.type>, !felt.type +// CHECK-LABEL: function.def @constrain(%arg0: !struct.type<@A<[@Start, @Stop, @Step]>>, %arg1: !felt.type) attributes {function.allow_constraint} {...}: +// CHECK-NEXT: <@A::@constrain>:function.return predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @constrain(%arg0: !struct.type<@A<[@Start, @Stop, @Step]>>, %arg1: !felt.type) attributes {function.allow_constraint} {...} + +// ----- + +module attributes {llzk.lang = "llzk"} { + struct.def @A { + struct.member @sums : !array.type<1 x !felt.type> + function.def @compute(%a: !felt.type, %b: !felt.type, %c: i1) -> !struct.type<@A> { + %c0 = arith.constant 0 : index + %self = struct.new : !struct.type<@A> + %sums = array.new : !array.type<1 x !felt.type> + + %ans = scf.if %c -> !felt.type { + %sum = felt.add %a, %b + scf.yield %sum : !felt.type + } else { + %diff = felt.add %a, %a + scf.yield %diff : !felt.type + } + %to_write = felt.sub %ans, %b + array.write %sums[%c0] = %to_write : !array.type<1 x !felt.type>, !felt.type + + struct.writem %self[@sums] = %sums : !struct.type<@A>, !array.type<1 x !felt.type> + function.return %self : !struct.type<@A> + } + + function.def @constrain(%self: !struct.type<@A>, %a: !felt.type, %b: !felt.type, %c: i1) { + function.return + } + } +} + +// CHECK-LABEL: function.def @compute(%arg0: !felt.type, %arg1: !felt.type, %arg2: i1) -> !struct.type<@A> attributes {function.allow_witness} {...}: +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} predecessors: +// CHECK-NEXT: <@A::@compute>:%array = array.new : <1 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%1 = felt.sub %0, %arg1 : !felt.type, !felt.type predecessors: +// CHECK-DAG: <@A::@compute>:scf.yield %2 : !felt.type +// CHECK-DAG: <@A::@compute>:scf.yield %2 : !felt.type +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %1 : <1 x !felt.type>, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%1 = felt.sub %0, %arg1 : !felt.type, !felt.type +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@sums] = %array : <@A>, !array.type<1 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:array.write %array[%c0] = %1 : <1 x !felt.type>, !felt.type +// CHECK-NEXT: <@A::@compute>:function.return %self : !struct.type<@A> predecessors: +// CHECK-NEXT: <@A::@compute>:struct.writem %self[@sums] = %array : <@A>, !array.type<1 x !felt.type> +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @compute(%arg0: !felt.type, %arg1: !felt.type, %arg2: i1) -> !struct.type<@A> attributes {function.allow_witness} {...} +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A> predecessors: +// CHECK-NEXT: <@A::@compute>:%c0 = arith.constant 0 : index +// CHECK-NEXT: <@A::@compute>:%array = array.new : <1 x !felt.type> predecessors: +// CHECK-NEXT: <@A::@compute>:%self = struct.new : <@A> +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg1 : !felt.type, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} +// CHECK-NEXT: <@A::@compute>:scf.yield %2 : !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg1 : !felt.type, !felt.type +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg0 : !felt.type, !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%0 = scf.if %arg2 -> (!felt.type) {...} else {...} +// CHECK-NEXT: <@A::@compute>:scf.yield %2 : !felt.type predecessors: +// CHECK-NEXT: <@A::@compute>:%2 = felt.add %arg0, %arg0 : !felt.type, !felt.type +// CHECK-LABEL: function.def @constrain(%arg0: !struct.type<@A>, %arg1: !felt.type, %arg2: !felt.type, %arg3: i1) attributes {function.allow_constraint} {...}: +// CHECK-NEXT: <@A::@constrain>:function.return predecessors: +// CHECK-NEXT: <(no parent function op)>:function.def @constrain(%arg0: !struct.type<@A>, %arg1: !felt.type, %arg2: !felt.type, %arg3: i1) attributes {function.allow_constraint} {...} diff --git a/test/FrontendLang/Zirgen/Inputs/zir_example_loops.llzk b/test/FrontendLang/Zirgen/Inputs/zir_example_loops.llzk index ddb67470b..40fee7a72 100644 --- a/test/FrontendLang/Zirgen/Inputs/zir_example_loops.llzk +++ b/test/FrontendLang/Zirgen/Inputs/zir_example_loops.llzk @@ -389,7 +389,7 @@ module attributes {llzk.lang = "zirgen"} { struct.member @b0 : !struct.type<@NondetReg<[]>> {column} struct.member @"$temp_0" : !struct.type<@NondetU8Reg<[]>> struct.member @"$temp" : !felt.type - function.def @compute(%arg0: !struct.type<@ValU32<[]>>, %arg1: !felt.type) -> !struct.type<@ExpandU32<[]>> attributes {function.allow_witness} { + function.def @compute(%arg0: !struct.type<@ValU32<[]>>, %arg1: !felt.type) -> !struct.type<@ExpandU32<[]>> attributes {function.allow_witness, function.allow_non_native_field_ops} { %felt_const_2 = felt.const 2 %felt_const_32768 = felt.const 32768 %felt_const_128 = felt.const 128 @@ -715,7 +715,7 @@ module attributes {llzk.lang = "zirgen"} { struct.member @"$temp_0" : !struct.type<@Div<[]>> struct.member @"$temp" : !felt.type struct.member @inst : !struct.type<@ValU32<[]>> - function.def @compute(%arg0: !struct.type<@ValU32<[]>>) -> !struct.type<@Decoder<[]>> attributes {function.allow_witness} { + function.def @compute(%arg0: !struct.type<@ValU32<[]>>) -> !struct.type<@Decoder<[]>> attributes {function.allow_witness, function.allow_non_native_field_ops} { %felt_const_65520 = felt.const 65520 %felt_const_65535 = felt.const 65535 %felt_const_61440 = felt.const 61440