diff --git a/backends/pcl/lib/Conversion/PCLLoweringPass.cpp b/backends/pcl/lib/Conversion/PCLLoweringPass.cpp index ad28d3b267..8bb99244e2 100644 --- a/backends/pcl/lib/Conversion/PCLLoweringPass.cpp +++ b/backends/pcl/lib/Conversion/PCLLoweringPass.cpp @@ -657,11 +657,10 @@ static bool isStep1LegalOp(Operation *op) { static void populateStep1ConversionTarget(ConversionTarget &target, NonDetOpNames &names) { target.addLegalDialect(); target.addLegalOp(); - target.addDynamicallyLegalDialect< - BoolDialect, FeltDialect, CastDialect, arith::ArithDialect, ConstrainDialect, - array::ArrayDialect, global::GlobalDialect, include::IncludeDialect, pod::PODDialect, - polymorphic::PolymorphicDialect, ram::RAMDialect, smt::SMTDialect, string::StringDialect, - verif::VerifDialect, LLZKDialect, StructDialect, FunctionDialect>(isStep1LegalOp); + // Step 1 converts only struct constrain functions. Use the fallback for + // operations from any dialect that may appear elsewhere; the callback still + // rejects unconverted operations inside constrain functions. + target.markUnknownOpDynamicallyLegal(isStep1LegalOp); target.addDynamicallyLegalOp([&names](NonDetOp op) { return isStep1LegalOp(op) && names.find(op) == names.end(); diff --git a/changelogs/unreleased/krushimir__arithmetize-scf-if.yaml b/changelogs/unreleased/krushimir__arithmetize-scf-if.yaml new file mode 100644 index 0000000000..7c9723c7a7 --- /dev/null +++ b/changelogs/unreleased/krushimir__arithmetize-scf-if.yaml @@ -0,0 +1,4 @@ +added: + - "`llzk-arithmetize-scf-if` pass that replaces constraint-side `scf.if` operations with select polynomials and gated branch constraints" +fixed: + - "Allow operations outside struct constraint functions to pass through the first PCL conversion step" diff --git a/include/llzk/Transforms/LLZKTransformationPasses.td b/include/llzk/Transforms/LLZKTransformationPasses.td index a680c35ec0..e77d319995 100644 --- a/include/llzk/Transforms/LLZKTransformationPasses.td +++ b/include/llzk/Transforms/LLZKTransformationPasses.td @@ -127,6 +127,33 @@ def EnforceNoMemberOverwritePass : LLZKPass<"llzk-enforce-no-overwrite"> { }]; } +def ArithmetizeSCFIfPass : LLZKPass<"llzk-arithmetize-scf-if"> { + let summary = + "Replace constraint-side scf.if operations with select polynomials"; + let description = [{ + Rewrites `scf.if` operations inside constraint-bearing functions (those + marked `function.allow_constraint`) into straight-line felt arithmetic. + Felt results become + + else + condition * (then - else) + + where the `i1` condition is converted to a felt. Equality constraints in + the then and else regions are gated by `condition` and `1 - condition`, + respectively, so constraints from the inactive branch do not apply. + Nested conditionals accumulate the enclosing gates. + + Both branches execute after the rewrite. The pass therefore accepts only + speculatable operations (including `llzk.nondet`) and felt equality + constraints in the regions. Other effects and non-felt results or + constraints are rejected. + + Run `-llzk-algebraize-felt-ops` first when branch bodies contain felt + division or non-native felt operations. Algebraization turns their + branch-local validity constraints into equality constraints that this + pass can gate safely. + }]; +} + def WhileToForPass : LLZKPass<"llzk-while-to-for"> { let summary = "Converts scf.while loops to equivalent scf.for loops when possible"; diff --git a/lib/Transforms/LLZKArithmetizeSCFIfPass.cpp b/lib/Transforms/LLZKArithmetizeSCFIfPass.cpp new file mode 100644 index 0000000000..283ed6b675 --- /dev/null +++ b/lib/Transforms/LLZKArithmetizeSCFIfPass.cpp @@ -0,0 +1,288 @@ +//===-- LLZKArithmetizeSCFIfPass.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-arithmetize-scf-if` pass. +/// +//===----------------------------------------------------------------------===// + +#include "llzk/Dialect/Cast/IR/Dialect.h" +#include "llzk/Dialect/Cast/IR/Ops.h" +#include "llzk/Dialect/Constrain/IR/Dialect.h" +#include "llzk/Dialect/Constrain/IR/Ops.h" +#include "llzk/Dialect/Felt/IR/Attrs.h" +#include "llzk/Dialect/Felt/IR/Dialect.h" +#include "llzk/Dialect/Felt/IR/Ops.h" +#include "llzk/Dialect/Function/IR/OpTraits.h" +#include "llzk/Dialect/Function/IR/Ops.h" +#include "llzk/Transforms/LLZKTransformationPasses.h" + +#include +#include +#include + +#include +#include +#include +#include + +#include + +namespace llzk { +#define GEN_PASS_DEF_ARITHMETIZESCFIFPASS +#include "llzk/Transforms/LLZKTransformationPasses.h.inc" +} // namespace llzk + +#define DEBUG_TYPE "llzk-arithmetize-scf-if" + +using namespace mlir; +using namespace llzk; +using namespace llzk::cast; +using namespace llzk::constrain; +using namespace llzk::felt; +using namespace llzk::function; + +namespace { + +/// Rewrites every `scf.if` in one constraint-bearing function. The rewrite is +/// bottom-up so an outer branch can add its gate to constraints produced while +/// lowering nested conditionals. +class ConstraintIfArithmetizer { +public: + explicit ConstraintIfArithmetizer(FuncDefOp func) : fn(func), rewriter(func.getContext()) {} + + LogicalResult run(); + +private: + struct Gates { + Value whenTrue; + Value whenFalse; + }; + + FuncDefOp fn; + IRRewriter rewriter; + + /// Constants shared by the entire function. They live in an entry-block + /// prelude so they dominate conditionals in sibling regions. + llvm::DenseMap, Value> constPool; + Operation *lastPreludeOp = nullptr; + + LogicalResult validate(mlir::scf::IfOp ifOp); + LogicalResult validateBranch(mlir::scf::IfOp ifOp, Block &block); + bool isSafeToExecuteUnconditionally(Operation *op); + bool requiresFeltAlgebraization(Operation *op); + + Value feltConst(Location loc, FeltType type, uint64_t value); + Gates getGates(mlir::scf::IfOp ifOp, FeltType type, llvm::DenseMap &gates); + void gateConstraints( + mlir::scf::IfOp ifOp, Block &block, bool thenBranch, llvm::DenseMap &gates + ); + void inlineBranch(mlir::scf::IfOp ifOp, Block &block); + void rewrite(mlir::scf::IfOp ifOp); +}; + +bool ConstraintIfArithmetizer::requiresFeltAlgebraization(Operation *op) { + if (llvm::isa(op)) { + return true; + } + return op->getName().getDialectNamespace() != "bool" && op->hasTrait(); +} + +bool ConstraintIfArithmetizer::isSafeToExecuteUnconditionally(Operation *op) { + return mlir::isSpeculatable(op); +} + +LogicalResult ConstraintIfArithmetizer::validateBranch(mlir::scf::IfOp ifOp, Block &block) { + auto yield = llvm::dyn_cast(block.getTerminator()); + if (!yield) { + ifOp.emitError("expected each scf.if branch to end with scf.yield"); + return failure(); + } + + for (Operation &op : block.without_terminator()) { + if (llvm::isa(op)) { + continue; + } + if (auto eq = llvm::dyn_cast(op)) { + if (!llvm::isa(eq.getLhs().getType())) { + eq.emitError("scf.if arithmetization supports only felt equality constraints"); + return failure(); + } + continue; + } + if (requiresFeltAlgebraization(&op)) { + op.emitError("felt division and non-native operations must be lowered before scf.if"); + return failure(); + } + if (op.getNumRegions() != 0 || op.getNumSuccessors() != 0 || + !isSafeToExecuteUnconditionally(&op)) { + op.emitError() << "cannot arithmetize scf.if branch operation '" << op.getName() + << "' because executing both branches may change its effects"; + return failure(); + } + } + return success(); +} + +LogicalResult ConstraintIfArithmetizer::validate(mlir::scf::IfOp ifOp) { + if (!ifOp.getThenRegion().hasOneBlock() || + (!ifOp.getElseRegion().empty() && !ifOp.getElseRegion().hasOneBlock())) { + ifOp.emitError("scf.if arithmetization requires single-block branches"); + return failure(); + } + if (llvm::any_of(ifOp.getResultTypes(), [](Type type) { return !llvm::isa(type); })) { + ifOp.emitError("scf.if arithmetization supports only felt results"); + return failure(); + } + if (failed(validateBranch(ifOp, ifOp.getThenRegion().front()))) { + return failure(); + } + if (!ifOp.getElseRegion().empty() && failed(validateBranch(ifOp, ifOp.getElseRegion().front()))) { + return failure(); + } + return success(); +} + +Value ConstraintIfArithmetizer::feltConst(Location loc, FeltType type, uint64_t value) { + auto [it, inserted] = constPool.try_emplace({type, value}); + if (!inserted) { + return it->second; + } + + OpBuilder::InsertionGuard guard(rewriter); + if (lastPreludeOp) { + rewriter.setInsertionPointAfter(lastPreludeOp); + } else { + rewriter.setInsertionPointToStart(&fn.getBody().front()); + } + auto attr = FeltConstAttr::get(rewriter.getContext(), llvm::APInt(2, value), type); + it->second = rewriter.create(loc, type, attr).getResult(); + lastPreludeOp = it->second.getDefiningOp(); + return it->second; +} + +ConstraintIfArithmetizer::Gates ConstraintIfArithmetizer::getGates( + mlir::scf::IfOp ifOp, FeltType type, llvm::DenseMap &gates +) { + auto [it, inserted] = gates.try_emplace(type); + if (inserted) { + rewriter.setInsertionPoint(ifOp); + Value whenTrue = + rewriter.create(ifOp.getLoc(), type, ifOp.getCondition()).getResult(); + Value whenFalse = + rewriter.create(ifOp.getLoc(), feltConst(ifOp.getLoc(), type, 1), whenTrue); + it->second = {whenTrue, whenFalse}; + } + return it->second; +} + +void ConstraintIfArithmetizer::gateConstraints( + mlir::scf::IfOp ifOp, Block &block, bool thenBranch, llvm::DenseMap &gates +) { + for (Operation &op : llvm::make_early_inc_range(block.without_terminator())) { + auto eq = llvm::dyn_cast(op); + if (!eq) { + continue; + } + + auto type = llvm::cast(eq.getLhs().getType()); + Gates branchGates = getGates(ifOp, type, gates); + Value gate = thenBranch ? branchGates.whenTrue : branchGates.whenFalse; + + rewriter.setInsertionPoint(eq); + Value difference = rewriter.create(eq.getLoc(), eq.getLhs(), eq.getRhs()); + Value gated = rewriter.create(eq.getLoc(), gate, difference); + rewriter.create(eq.getLoc(), gated, feltConst(eq.getLoc(), type, 0)); + rewriter.eraseOp(eq); + } +} + +void ConstraintIfArithmetizer::inlineBranch(mlir::scf::IfOp ifOp, Block &block) { + rewriter.eraseOp(block.getTerminator()); + rewriter.inlineBlockBefore(&block, ifOp); +} + +void ConstraintIfArithmetizer::rewrite(mlir::scf::IfOp ifOp) { + llvm::DenseMap gates; + + Block &thenBlock = ifOp.getThenRegion().front(); + auto thenYield = llvm::cast(thenBlock.getTerminator()); + SmallVector thenValues(thenYield.getOperands()); + gateConstraints(ifOp, thenBlock, /*thenBranch=*/true, gates); + + SmallVector elseValues; + Block *elseBlock = nullptr; + if (!ifOp.getElseRegion().empty()) { + elseBlock = &ifOp.getElseRegion().front(); + auto elseYield = llvm::cast(elseBlock->getTerminator()); + elseValues.append(elseYield.getOperands().begin(), elseYield.getOperands().end()); + gateConstraints(ifOp, *elseBlock, /*thenBranch=*/false, gates); + } + + inlineBranch(ifOp, thenBlock); + if (elseBlock) { + inlineBranch(ifOp, *elseBlock); + } + + rewriter.setInsertionPoint(ifOp); + SmallVector replacements; + replacements.reserve(ifOp.getNumResults()); + for (auto [type, thenValue, elseValue] : + llvm::zip_equal(ifOp.getResultTypes(), thenValues, elseValues)) { + Gates resultGates = getGates(ifOp, llvm::cast(type), gates); + Value difference = rewriter.create(ifOp.getLoc(), thenValue, elseValue); + Value selected = rewriter.create(ifOp.getLoc(), resultGates.whenTrue, difference); + replacements.push_back(rewriter.create(ifOp.getLoc(), elseValue, selected)); + } + rewriter.replaceOp(ifOp, replacements); +} + +LogicalResult ConstraintIfArithmetizer::run() { + SmallVector ifOps; + fn.walk([&](mlir::scf::IfOp ifOp) { ifOps.push_back(ifOp); }); + + // Validate the whole function before changing it. Nested scf.if operations + // are allowed here because the rewrite list is already bottom-up. + for (mlir::scf::IfOp ifOp : ifOps) { + if (failed(validate(ifOp))) { + return failure(); + } + } + for (mlir::scf::IfOp ifOp : ifOps) { + rewrite(ifOp); + } + return success(); +} + +class PassImpl : public llzk::impl::ArithmetizeSCFIfPassBase { + using Base = ArithmetizeSCFIfPassBase; + using Base::Base; + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert< + llzk::cast::CastDialect, llzk::constrain::ConstrainDialect, llzk::felt::FeltDialect>(); + } + + void runOnOperation() override { + SmallVector funcs; + getOperation()->walk([&](FuncDefOp fn) { + if (fn.hasAllowConstraintAttr() && !fn.getBody().empty()) { + funcs.push_back(fn); + } + }); + + for (FuncDefOp fn : funcs) { + if (failed(ConstraintIfArithmetizer(fn).run())) { + signalPassFailure(); + return; + } + } + } +}; +} // namespace diff --git a/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_fail.llzk b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_fail.llzk new file mode 100644 index 0000000000..1214073d4f --- /dev/null +++ b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_fail.llzk @@ -0,0 +1,87 @@ +// RUN: llzk-opt -split-input-file -llzk-arithmetize-scf-if -verify-diagnostics %s + +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + function.def @callee() { + function.return + } + struct.def @CallEffect { + function.def @compute(%cond: i1) -> !struct.type<@CallEffect> { + %self = struct.new : <@CallEffect> + function.return %self : !struct.type<@CallEffect> + } + function.def @constrain(%self: !struct.type<@CallEffect>, %cond: i1) attributes {function.allow_constraint} { + scf.if %cond { + // expected-error@+1 {{cannot arithmetize scf.if branch operation 'function.call' because executing both branches may change its effects}} + function.call @callee() : () -> () + } + function.return + } + } +} + +// ----- + +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @IndexResult { + function.def @compute(%cond: i1) -> !struct.type<@IndexResult> { + %self = struct.new : <@IndexResult> + function.return %self : !struct.type<@IndexResult> + } + function.def @constrain(%self: !struct.type<@IndexResult>, %cond: i1) attributes {function.allow_constraint} { + // expected-error@+1 {{scf.if arithmetization supports only felt results}} + %result = scf.if %cond -> index { + %one = arith.constant 1 : index + scf.yield %one : index + } else { + %zero = arith.constant 0 : index + scf.yield %zero : index + } + function.return + } + } +} + +// ----- + +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @IntegerConstraint { + function.def @compute(%cond: i1) -> !struct.type<@IntegerConstraint> { + %self = struct.new : <@IntegerConstraint> + function.return %self : !struct.type<@IntegerConstraint> + } + function.def @constrain(%self: !struct.type<@IntegerConstraint>, %cond: i1) attributes {function.allow_constraint} { + scf.if %cond { + %one = arith.constant 1 : i1 + // expected-error@+1 {{scf.if arithmetization supports only felt equality constraints}} + constrain.eq %cond, %one : i1, i1 + } + function.return + } + } +} + +// ----- + +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @Division { + function.def @compute(%a: !F, %b: !F, %cond: i1) -> !struct.type<@Division> { + %self = struct.new : <@Division> + function.return %self : !struct.type<@Division> + } + function.def @constrain(%self: !struct.type<@Division>, %a: !F, %b: !F, %cond: i1) attributes {function.allow_constraint} { + scf.if %cond { + // expected-error@+1 {{felt division and non-native operations must be lowered before scf.if}} + %quotient = felt.div %a, %b : !F, !F + } + function.return + } + } +} diff --git a/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pass.llzk b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pass.llzk new file mode 100644 index 0000000000..f3685af243 --- /dev/null +++ b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pass.llzk @@ -0,0 +1,136 @@ +// NOTE: Assertions have been autogenerated by scripts/generate-test-checks.py + +// The script is designed to make adding checks to +// a test case fast, it is *not* designed to be authoritative +// about what constitutes a good test! The CHECK should be +// minimized and named to reflect the test intent. + +// RUN: llzk-opt -split-input-file -llzk-arithmetize-scf-if %s | FileCheck --enable-var-scope %s + +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @Select { +// CHECK-NEXT: struct.member @out : !felt.type<"koalabear"> {llzk.pub} +// CHECK-NEXT: function.def @compute(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">, %[[VAL_1:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">) -> !struct.type<@Select> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@Select> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@Select> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@Select>, %[[VAL_4:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">, %[[VAL_5:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = felt.const 1 : <"koalabear"> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = felt.const 0 : <"koalabear"> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = bool.cmp lt(%[[VAL_4]], %[[VAL_5]]) : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_8]] : i1, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_6]], %[[VAL_9]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_4]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_11]], %[[VAL_4]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_9]], %[[VAL_12]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_13]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_4]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_15:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_14]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_10]], %[[VAL_15]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_16]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_17:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_11]], %[[VAL_14]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_18:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_9]], %[[VAL_17]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_19:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_14]], %[[VAL_18]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_20:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_3]][@out] : <@Select>, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_19]], %[[VAL_20]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @Select { + struct.member @out : !F {llzk.pub} + function.def @compute(%a: !F, %b: !F) -> !struct.type<@Select> { + %self = struct.new : <@Select> + function.return %self : !struct.type<@Select> + } + function.def @constrain(%self: !struct.type<@Select>, %a: !F, %b: !F) attributes {function.allow_constraint} { + %cond = bool.cmp lt(%a, %b) : !F, !F + %selected = scf.if %cond -> !F { + %then = felt.add %a, %b : !F, !F + constrain.eq %then, %a : !F, !F + scf.yield %then : !F + } else { + %else = felt.sub %a, %b : !F, !F + constrain.eq %else, %b : !F, !F + scf.yield %else : !F + } + %out = struct.readm %self[@out] : !struct.type<@Select>, !F + constrain.eq %selected, %out : !F, !F + function.return + } + } +} + +// ----- + +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @Nested { +// CHECK-NEXT: struct.member @out : !felt.type<"koalabear"> {llzk.pub} +// CHECK-NEXT: function.def @compute(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">, %[[VAL_1:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">) -> !struct.type<@Nested> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@Nested> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@Nested> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@Nested>, %[[VAL_4:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">, %[[VAL_5:[0-9a-zA-Z_\.]+]]: !felt.type<"koalabear">) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = felt.const 1 : <"koalabear"> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = felt.const 0 : <"koalabear"> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = bool.cmp lt(%[[VAL_4]], %[[VAL_5]]) : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = bool.cmp eq(%[[VAL_4]], %[[VAL_5]]) : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_8]] : i1, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_6]], %[[VAL_10]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_9]] : i1, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_6]], %[[VAL_12]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_4]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_15:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_12]], %[[VAL_14]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_15]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_17:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_10]], %[[VAL_16]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_17]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_18:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_5]], %[[VAL_4]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_19:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_13]], %[[VAL_18]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_20:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_19]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_21:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_10]], %[[VAL_20]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_21]], %[[VAL_7]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_22:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_4]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_23:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_12]], %[[VAL_22]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_24:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_5]], %[[VAL_23]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_25:[0-9a-zA-Z_\.]+]] = felt.sub %[[VAL_24]], %[[VAL_5]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_26:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_10]], %[[VAL_25]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_27:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_5]], %[[VAL_26]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: %[[VAL_28:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_3]][@out] : <@Nested>, !felt.type<"koalabear"> +// CHECK-NEXT: constrain.eq %[[VAL_27]], %[[VAL_28]] : !felt.type<"koalabear">, !felt.type<"koalabear"> +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @Nested { + struct.member @out : !F {llzk.pub} + function.def @compute(%a: !F, %b: !F) -> !struct.type<@Nested> { + %self = struct.new : <@Nested> + function.return %self : !struct.type<@Nested> + } + function.def @constrain(%self: !struct.type<@Nested>, %a: !F, %b: !F) attributes {function.allow_constraint} { + %outer_cond = bool.cmp lt(%a, %b) : !F, !F + %inner_cond = bool.cmp eq(%a, %b) : !F, !F + %selected = scf.if %outer_cond -> !F { + %inner = scf.if %inner_cond -> !F { + constrain.eq %a, %b : !F, !F + scf.yield %a : !F + } else { + constrain.eq %b, %a : !F, !F + scf.yield %b : !F + } + scf.yield %inner : !F + } else { + scf.yield %b : !F + } + %out = struct.readm %self[@out] : !struct.type<@Nested>, !F + constrain.eq %selected, %out : !F, !F + function.return + } + } +} diff --git a/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pcl_e2e.llzk b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pcl_e2e.llzk new file mode 100644 index 0000000000..72951278e6 --- /dev/null +++ b/test/Transforms/ArithmetizeSCFIf/arithmetize_scf_if_pcl_e2e.llzk @@ -0,0 +1,62 @@ +// NOTE: Assertions have been autogenerated by scripts/generate-test-checks.py + +// The script is designed to make adding checks to +// a test case fast, it is *not* designed to be authoritative +// about what constitutes a good test! The CHECK should be +// minimized and named to reflect the test intent. + +// REQUIRES: with-pcl +// RUN: llzk-opt -llzk-arithmetize-scf-if -llzk-to-pcl %s | FileCheck %s + +!F = !felt.type<"koalabear"> + +// CHECK-LABEL: module attributes {pcl.prime = #pcl.prime<2130706433 : i65>} { +// CHECK-NEXT: func.func @Main(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !pcl.felt, %[[VAL_1:[0-9a-zA-Z_\.]+]]: !pcl.felt) -> !pcl.felt { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = pcl.const 1 : i2 +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = pcl.const 0 : i2 +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = pcl.lt %[[VAL_0]], %[[VAL_1]] +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = pcl.asfelt %[[VAL_4]] +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = pcl.sub %[[VAL_2]], %[[VAL_5]] +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = pcl.add %[[VAL_0]], %[[VAL_1]] +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = pcl.sub %[[VAL_7]], %[[VAL_0]] +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = pcl.mul %[[VAL_5]], %[[VAL_8]] +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = pcl.eq %[[VAL_9]], %[[VAL_3]] +// CHECK-NEXT: pcl.assert %[[VAL_10]] +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = pcl.sub %[[VAL_0]], %[[VAL_1]] +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = pcl.sub %[[VAL_11]], %[[VAL_1]] +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = pcl.mul %[[VAL_6]], %[[VAL_12]] +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = pcl.eq %[[VAL_13]], %[[VAL_3]] +// CHECK-NEXT: pcl.assert %[[VAL_14]] +// CHECK-NEXT: %[[VAL_15:[0-9a-zA-Z_\.]+]] = pcl.sub %[[VAL_7]], %[[VAL_11]] +// CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = pcl.mul %[[VAL_5]], %[[VAL_15]] +// CHECK-NEXT: %[[VAL_17:[0-9a-zA-Z_\.]+]] = pcl.add %[[VAL_11]], %[[VAL_16]] +// CHECK-NEXT: %[[VAL_18:[0-9a-zA-Z_\.]+]] = pcl.var "out" true +// CHECK-NEXT: %[[VAL_19:[0-9a-zA-Z_\.]+]] = pcl.eq %[[VAL_17]], %[[VAL_18]] +// CHECK-NEXT: pcl.assert %[[VAL_19]] +// CHECK-NEXT: return %[[VAL_18]] : !pcl.felt +// CHECK-NEXT: } +// CHECK-NEXT: } +module attributes {llzk.lang} { + struct.def @Main { + struct.member @out : !F {llzk.pub} + function.def @compute(%a: !F, %b: !F) -> !struct.type<@Main> { + %self = struct.new : <@Main> + function.return %self : !struct.type<@Main> + } + function.def @constrain(%self: !struct.type<@Main>, %a: !F, %b: !F) attributes {function.allow_constraint} { + %cond = bool.cmp lt(%a, %b) : !F, !F + %selected = scf.if %cond -> !F { + %then = felt.add %a, %b : !F, !F + constrain.eq %then, %a : !F, !F + scf.yield %then : !F + } else { + %else = felt.sub %a, %b : !F, !F + constrain.eq %else, %b : !F, !F + scf.yield %else : !F + } + %out = struct.readm %self[@out] : !struct.type<@Main>, !F + constrain.eq %selected, %out : !F, !F + function.return + } + } +} diff --git a/test/Transforms/PCLLowering/pcl_compute_body_legality.llzk b/test/Transforms/PCLLowering/pcl_compute_body_legality.llzk new file mode 100644 index 0000000000..827e4e4839 --- /dev/null +++ b/test/Transforms/PCLLowering/pcl_compute_body_legality.llzk @@ -0,0 +1,42 @@ +// NOTE: Assertions have been autogenerated by scripts/generate-test-checks.py + +// The script is designed to make adding checks to +// a test case fast, it is *not* designed to be authoritative +// about what constitutes a good test! The CHECK should be +// minimized and named to reflect the test intent. + +// REQUIRES: with-pcl +// RUN: llzk-opt -llzk-to-pcl %s | FileCheck %s + +// Step 1 lowers only struct constraint functions. Operations in compute +// functions, including dialects not named by the conversion, must remain legal +// until the containing struct is lowered. + +!F = !felt.type<"koalabear"> + +// CHECK-LABEL: module attributes {pcl.prime = #pcl.prime<2130706433 : i65>} { +// CHECK-NEXT: func.func @ComputeSCF(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !pcl.felt) -> !pcl.felt { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = pcl.var "out" true +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = pcl.eq %[[VAL_1]], %[[VAL_0]] +// CHECK-NEXT: pcl.assert %[[VAL_2]] +// CHECK-NEXT: return %[[VAL_1]] : !pcl.felt +// CHECK-NEXT: } +// CHECK-NEXT: } +module attributes {llzk.lang} { + struct.def @ComputeSCF { + struct.member @out : !F {llzk.pub} + function.def @compute(%a: !F) -> !struct.type<@ComputeSCF> { + %self = struct.new : <@ComputeSCF> + %condition = arith.constant true + scf.if %condition { + struct.writem %self[@out] = %a : !struct.type<@ComputeSCF>, !F + } + function.return %self : !struct.type<@ComputeSCF> + } + function.def @constrain(%self: !struct.type<@ComputeSCF>, %a: !F) { + %out = struct.readm %self[@out] : !struct.type<@ComputeSCF>, !F + constrain.eq %out, %a : !F, !F + function.return + } + } +} diff --git a/test/Transforms/PCLLowering/pcl_compute_body_legality_fail.llzk b/test/Transforms/PCLLowering/pcl_compute_body_legality_fail.llzk new file mode 100644 index 0000000000..9affd91914 --- /dev/null +++ b/test/Transforms/PCLLowering/pcl_compute_body_legality_fail.llzk @@ -0,0 +1,22 @@ +// REQUIRES: with-pcl +// RUN: llzk-opt -llzk-to-pcl -verify-diagnostics %s + +// The fallback legality rule applies only outside struct constraint functions. + +!F = !felt.type<"koalabear"> + +module attributes {llzk.lang} { + struct.def @ConstraintSCF { + function.def @compute(%a: !F, %b: !F) -> !struct.type<@ConstraintSCF> { + %self = struct.new : <@ConstraintSCF> + function.return %self : !struct.type<@ConstraintSCF> + } + function.def @constrain(%self: !struct.type<@ConstraintSCF>, %a: !F, %b: !F) { + %condition = bool.cmp eq(%a, %b) : !F, !F + // expected-error@+1 {{failed to legalize operation 'scf.if'}} + scf.if %condition { + } + function.return + } + } +}