From d873ee47febb90055ba7a8e6ff54bcb8402763b9 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:20:39 +0800 Subject: [PATCH 01/12] Emit typed auxiliary members in lowering passes Fixes project-llzk/llzk-lib#552. --- .../lib/r1cs/Transforms/R1CSLoweringPass.cpp | 4 +- .../unreleased/poly-typed-aux-members.yaml | 3 ++ include/llzk/Transforms/LLZKLoweringUtils.h | 4 ++ lib/Transforms/LLZKLoweringUtils.cpp | 13 +++++-- lib/Transforms/LLZKPolyLoweringPass.cpp | 6 +-- .../poly_lowering_typed_aux_member.llzk | 39 +++++++++++++++++++ .../r1cs_lowering_typed_aux_member.llzk | 34 ++++++++++++++++ 7 files changed, 94 insertions(+), 9 deletions(-) create mode 100644 changelogs/unreleased/poly-typed-aux-members.yaml create mode 100644 test/Transforms/PolyLowering/poly_lowering_typed_aux_member.llzk create mode 100644 test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk diff --git a/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp b/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp index 87ba831ed..cbd88787a 100644 --- a/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp +++ b/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp @@ -340,7 +340,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase { if (degLhs == 2 && degRhs == 2) { builder.setInsertionPoint(op); std::string auxName = R1CS_AUXILIARY_MEMBER_PREFIX + std::to_string(auxCounter++); - MemberDefOp auxMember = addAuxMember(structDef, auxName); + MemberDefOp auxMember = addAuxMember(structDef, auxName, val.getType()); Value aux = builder.create( val.getLoc(), val.getType(), constrainFunc.getSelfValueFromConstrain(), auxMember.getNameAttr() @@ -673,7 +673,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase { if (degLhs == 2 && degRhs == 2) { builder.setInsertionPoint(eqOp); std::string auxName = R1CS_AUXILIARY_MEMBER_PREFIX + std::to_string(auxCounter++); - MemberDefOp auxMember = addAuxMember(structDef, auxName); + MemberDefOp auxMember = addAuxMember(structDef, auxName, lhs.getType()); Value aux = builder.create( eqOp.getLoc(), lhs.getType(), constrainFunc.getSelfValueFromConstrain(), auxMember.getNameAttr() diff --git a/changelogs/unreleased/poly-typed-aux-members.yaml b/changelogs/unreleased/poly-typed-aux-members.yaml new file mode 100644 index 000000000..c73ef93c5 --- /dev/null +++ b/changelogs/unreleased/poly-typed-aux-members.yaml @@ -0,0 +1,3 @@ +fixed: + - Emit polynomial-lowering auxiliary members with the exact type of the materialized expression + - Emit all R1CS-lowering auxiliary members with the exact type of the materialized expression diff --git a/include/llzk/Transforms/LLZKLoweringUtils.h b/include/llzk/Transforms/LLZKLoweringUtils.h index 3a346dbd2..8ccaee945 100644 --- a/include/llzk/Transforms/LLZKLoweringUtils.h +++ b/include/llzk/Transforms/LLZKLoweringUtils.h @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -39,6 +40,9 @@ checkForAuxMemberConflicts(component::StructDefOp structDef, llvm::StringRef aux mlir::LogicalResult checkConstrainBodyIsStraightLine(function::FuncDefOp constrainFunc, llvm::StringRef passName); +component::MemberDefOp +addAuxMember(component::StructDefOp structDef, llvm::StringRef name, mlir::Type type); + component::MemberDefOp addAuxMember(component::StructDefOp structDef, llvm::StringRef name); unsigned getFeltDegree(mlir::Value val, llvm::DenseMap &memo); diff --git a/lib/Transforms/LLZKLoweringUtils.cpp b/lib/Transforms/LLZKLoweringUtils.cpp index 34effc94d..15401eff6 100644 --- a/lib/Transforms/LLZKLoweringUtils.cpp +++ b/lib/Transforms/LLZKLoweringUtils.cpp @@ -148,12 +148,17 @@ void replaceSubsequentUsesWith(Value oldVal, Value newVal, Operation *afterOp) { } } -MemberDefOp addAuxMember(StructDefOp structDef, StringRef name) { +MemberDefOp addAuxMember(StructDefOp structDef, StringRef name, Type type) { + assert(type && "auxiliary member type must be non-null"); + OpBuilder builder(structDef); builder.setInsertionPointToEnd(structDef.getBody()); - return builder.create( - structDef.getLoc(), builder.getStringAttr(name), builder.getType() - ); + return builder.create(structDef.getLoc(), builder.getStringAttr(name), type); +} + +MemberDefOp addAuxMember(StructDefOp structDef, StringRef name) { + OpBuilder builder(structDef); + return addAuxMember(structDef, name, builder.getType()); } unsigned getFeltDegree(Value val, DenseMap &memo) { diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index 2784686fe..ca8e84d49 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -178,7 +178,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { // Optimization: If lhs == rhs, factor it only once if (lhs == rhs && eraseMul) { std::string auxName = AUXILIARY_MEMBER_PREFIX + std::to_string(this->auxCounter++); - MemberDefOp auxMember = addAuxMember(structDef, auxName); + MemberDefOp auxMember = addAuxMember(structDef, auxName, lhs.getType()); auto auxVal = builder.create( lhs.getLoc(), lhs.getType(), selfVal, auxMember.getNameAttr() @@ -206,7 +206,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { // Create auxiliary member for toFactor std::string auxName = AUXILIARY_MEMBER_PREFIX + std::to_string(this->auxCounter++); - MemberDefOp auxMember = addAuxMember(structDef, auxName); + MemberDefOp auxMember = addAuxMember(structDef, auxName, toFactor.getType()); // Read back as MemberReadOp (new SSA value) auto auxVal = builder.create( @@ -265,7 +265,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { // Callees only receive SSA values, not the caller expression tree, so nonlinear // call arguments must be represented by an auxiliary member read. std::string auxName = AUXILIARY_MEMBER_PREFIX + std::to_string(this->auxCounter++); - MemberDefOp auxMember = addAuxMember(structDef, auxName); + MemberDefOp auxMember = addAuxMember(structDef, auxName, loweredVal.getType()); OpBuilder builder(callOp); Value selfVal = constrainFunc.getSelfValueFromConstrain(); diff --git a/test/Transforms/PolyLowering/poly_lowering_typed_aux_member.llzk b/test/Transforms/PolyLowering/poly_lowering_typed_aux_member.llzk new file mode 100644 index 000000000..8d3dc05b3 --- /dev/null +++ b/test/Transforms/PolyLowering/poly_lowering_typed_aux_member.llzk @@ -0,0 +1,39 @@ +// RUN: llzk-opt -I %S -split-input-file -llzk-poly-lowering-pass="max-degree=2" --verify-each %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + struct.def @TypedAux { + struct.member @out : !felt.type<"babybear"> + + function.def @compute( + %a: !felt.type<"babybear">, + %b: !felt.type<"babybear">, + %c: !felt.type<"babybear"> + ) -> !struct.type<@TypedAux> { + %self = struct.new : !struct.type<@TypedAux> + function.return %self : !struct.type<@TypedAux> + } + + function.def @constrain( + %self: !struct.type<@TypedAux>, + %a: !felt.type<"babybear">, + %b: !felt.type<"babybear">, + %c: !felt.type<"babybear"> + ) { + %ab = felt.mul %a, %b : !felt.type<"babybear">, !felt.type<"babybear"> + %abc = felt.mul %ab, %c : !felt.type<"babybear">, !felt.type<"babybear"> + %out = struct.readm %self[@out] : !struct.type<@TypedAux>, !felt.type<"babybear"> + constrain.eq %out, %abc : !felt.type<"babybear"> + function.return + } + } +} + +// CHECK-LABEL: struct.def @TypedAux +// CHECK-LABEL: function.def @compute +// CHECK: %[[AB:.*]] = felt.mul %{{.*}}, %{{.*}} : !felt.type<"babybear">, !felt.type<"babybear"> +// CHECK: struct.writem %{{.*}}[@__llzk_poly_lowering_pass_aux_member_0] = %[[AB]] : <@TypedAux>, !felt.type<"babybear"> +// CHECK-LABEL: function.def @constrain +// CHECK: %[[AB_CONSTRAIN:.*]] = felt.mul %{{.*}}, %{{.*}} : !felt.type<"babybear">, !felt.type<"babybear"> +// CHECK: %[[AUX:.*]] = struct.readm %{{.*}}[@__llzk_poly_lowering_pass_aux_member_0] : <@TypedAux>, !felt.type<"babybear"> +// CHECK: constrain.eq %[[AUX]], %[[AB_CONSTRAIN]] : !felt.type<"babybear">, !felt.type<"babybear"> +// CHECK: struct.member @__llzk_poly_lowering_pass_aux_member_0 : !felt.type<"babybear"> diff --git a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk new file mode 100644 index 000000000..fc62e3f09 --- /dev/null +++ b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk @@ -0,0 +1,34 @@ +// RUN: llzk-opt -split-input-file -llzk-full-r1cs-lowering --verify-each %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + struct.def @TypedR1CSAux { + struct.member @out : !felt.type<"babybear"> {llzk.pub} + + function.def @compute( + %a: !felt.type<"babybear">, + %b: !felt.type<"babybear">, + %c: !felt.type<"babybear"> + ) -> !struct.type<@TypedR1CSAux> { + %self = struct.new : !struct.type<@TypedR1CSAux> + function.return %self : !struct.type<@TypedR1CSAux> + } + + function.def @constrain( + %self: !struct.type<@TypedR1CSAux>, + %a: !felt.type<"babybear"> {llzk.pub}, + %b: !felt.type<"babybear">, + %c: !felt.type<"babybear"> + ) { + %ab = felt.mul %a, %b : !felt.type<"babybear">, !felt.type<"babybear"> + %bc = felt.mul %b, %c : !felt.type<"babybear">, !felt.type<"babybear"> + constrain.eq %ab, %bc : !felt.type<"babybear"> + function.return + } + } +} + +// CHECK-LABEL: r1cs.circuit @TypedR1CSAux inputs +// CHECK-SAME: %{{.*}}: !r1cs.signal {#r1cs.pub} +// CHECK-SAME: %{{.*}}: !r1cs.signal +// CHECK-SAME: %{{.*}}: !r1cs.signal +// CHECK: r1cs.constrain From 3e2a14261a6a29db54cbb4cb4fd308f141baa941 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:22:02 +0800 Subject: [PATCH 02/12] Reuse poly rewrites only when they dominate the use Fixes project-llzk/llzk-lib#553. --- .../unreleased/poly-rewrite-dominance.yaml | 2 + lib/Transforms/LLZKPolyLoweringPass.cpp | 66 ++++++++++++++----- .../poly_lowering_rewrite_dominance.llzk | 45 +++++++++++++ 3 files changed, 98 insertions(+), 15 deletions(-) create mode 100644 changelogs/unreleased/poly-rewrite-dominance.yaml create mode 100644 test/Transforms/PolyLowering/poly_lowering_rewrite_dominance.llzk diff --git a/changelogs/unreleased/poly-rewrite-dominance.yaml b/changelogs/unreleased/poly-rewrite-dominance.yaml new file mode 100644 index 000000000..286b82f1d --- /dev/null +++ b/changelogs/unreleased/poly-rewrite-dominance.yaml @@ -0,0 +1,2 @@ +fixed: + - Avoid reusing polynomial-lowering replacement values at uses they do not dominate diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index ca8e84d49..02779e9a9 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -66,6 +66,22 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { }); } + static bool valueDominatesUse(Value value, Operation *useOp) { + if (!useOp || llvm::isa(value)) { + return true; + } + + Operation *defOp = value.getDefiningOp(); + if (!defOp) { + return true; + } + + // checkConstrainBodyIsStraightLine has already rejected multi-block or + // nested-region constrain bodies. Within that invariant, same-block order + // is a sufficient and exact dominance check for values introduced here. + return defOp->getBlock() == useOp->getBlock() && defOp->isBeforeInBlock(useOp); + } + // Recursively compute degree of FeltOps SSA values unsigned getDegree(Value val, DenseMap &memo) { if (auto it = memo.find(val); it != memo.end()) { @@ -105,27 +121,39 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { } Value lowerExpression( - Value val, StructDefOp structDef, FuncDefOp constrainFunc, + Value val, StructDefOp structDef, FuncDefOp constrainFunc, Operation *useOp, DenseMap °reeMemo, DenseMap &rewrites, SmallVector &auxAssignments ) { - if (rewrites.count(val)) { - return rewrites[val]; + auto rewriteIt = rewrites.find(val); + if (rewriteIt != rewrites.end() && valueDominatesUse(rewriteIt->second, useOp)) { + return rewriteIt->second; } + auto cacheIdentityRewriteIfAbsent = [&]() { + if (!rewrites.contains(val)) { + rewrites[val] = val; + } + }; + unsigned degree = getDegree(val, degreeMemo); if (degree <= maxDegree) { - rewrites[val] = val; + // A cached replacement that does not dominate this use may still be the + // right replacement for later uses. Return the original value for this + // use without clobbering that scoped rewrite. + cacheIdentityRewriteIfAbsent(); return val; } // Degree-neutral roots can still contain over-degree operands. auto lowerBinaryRoot = [&](auto op) -> Value { Value lhs = lowerExpression( - op.getLhs(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + op.getLhs(), structDef, constrainFunc, op.getOperation(), degreeMemo, rewrites, + auxAssignments ); Value rhs = lowerExpression( - op.getRhs(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + op.getRhs(), structDef, constrainFunc, op.getOperation(), degreeMemo, rewrites, + auxAssignments ); if (lhs != op.getLhs()) { @@ -135,7 +163,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { op.getRhsMutable().set(rhs); } degreeMemo[val] = std::max(getDegree(lhs, degreeMemo), getDegree(rhs, degreeMemo)); - rewrites[val] = val; + cacheIdentityRewriteIfAbsent(); return val; }; @@ -149,24 +177,27 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { if (auto negOp = val.getDefiningOp()) { Value operand = lowerExpression( - negOp.getOperand(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + negOp.getOperand(), structDef, constrainFunc, negOp.getOperation(), degreeMemo, rewrites, + auxAssignments ); if (operand != negOp.getOperand()) { negOp.getOperandMutable().set(operand); } degreeMemo[val] = getDegree(operand, degreeMemo); - rewrites[val] = val; + cacheIdentityRewriteIfAbsent(); return val; } if (auto mulOp = val.getDefiningOp()) { // Recursively lower operands first Value lhs = lowerExpression( - mulOp.getLhs(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + mulOp.getLhs(), structDef, constrainFunc, mulOp.getOperation(), degreeMemo, rewrites, + auxAssignments ); Value rhs = lowerExpression( - mulOp.getRhs(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + mulOp.getRhs(), structDef, constrainFunc, mulOp.getOperation(), degreeMemo, rewrites, + auxAssignments ); unsigned lhsDeg = getDegree(lhs, degreeMemo); @@ -246,7 +277,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { } // Unsupported roots are left unchanged. - rewrites[val] = val; + cacheIdentityRewriteIfAbsent(); return val; } @@ -256,7 +287,10 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { SmallVector &auxAssignments ) { Value loweredVal = - lowerExpression(val, structDef, constrainFunc, degreeMemo, rewrites, auxAssignments); + lowerExpression( + val, structDef, constrainFunc, callOp.getOperation(), degreeMemo, rewrites, + auxAssignments + ); DenseMap checkMemo; if (getDegree(loweredVal, checkMemo) <= 1) { return loweredVal; @@ -389,13 +423,15 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { if (degreeLhs > maxDegree) { Value loweredExpr = lowerExpression( - lhsOperand.get(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + lhsOperand.get(), structDef, constrainFunc, constraintOp.getOperation(), degreeMemo, + rewrites, auxAssignments ); lhsOperand.set(loweredExpr); } if (degreeRhs > maxDegree) { Value loweredExpr = lowerExpression( - rhsOperand.get(), structDef, constrainFunc, degreeMemo, rewrites, auxAssignments + rhsOperand.get(), structDef, constrainFunc, constraintOp.getOperation(), degreeMemo, + rewrites, auxAssignments ); rhsOperand.set(loweredExpr); } diff --git a/test/Transforms/PolyLowering/poly_lowering_rewrite_dominance.llzk b/test/Transforms/PolyLowering/poly_lowering_rewrite_dominance.llzk new file mode 100644 index 000000000..74b224ece --- /dev/null +++ b/test/Transforms/PolyLowering/poly_lowering_rewrite_dominance.llzk @@ -0,0 +1,45 @@ +// RUN: llzk-opt -I %S -split-input-file -llzk-poly-lowering-pass="max-degree=2" --verify-each %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang = "llzk"} { + struct.def @Child { + function.def @compute(%x: !felt.type) -> !struct.type<@Child> { + %self = struct.new : !struct.type<@Child> + function.return %self : !struct.type<@Child> + } + + function.def @constrain(%self: !struct.type<@Child>, %x: !felt.type) { + function.return + } + } + + struct.def @Parent { + struct.member @child : !struct.type<@Child> + struct.member @out : !felt.type {llzk.pub} + + function.def @compute(%a: !felt.type, %b: !felt.type) -> !struct.type<@Parent> { + %self = struct.new : !struct.type<@Parent> + %ab = felt.mul %a, %b : !felt.type, !felt.type + %cubeish = felt.mul %ab, %a : !felt.type, !felt.type + struct.writem %self[@out] = %cubeish : !struct.type<@Parent>, !felt.type + function.return %self : !struct.type<@Parent> + } + + function.def @constrain(%self: !struct.type<@Parent>, %a: !felt.type, %b: !felt.type) { + %ab = felt.mul %a, %b : !felt.type, !felt.type + %child = struct.readm %self[@child] : !struct.type<@Parent>, !struct.type<@Child> + function.call @Child::@constrain(%child, %ab) : (!struct.type<@Child>, !felt.type) -> () + %cubeish = felt.mul %ab, %a : !felt.type, !felt.type + %out = struct.readm %self[@out] : !struct.type<@Parent>, !felt.type + constrain.eq %out, %cubeish : !felt.type + function.return + } + } +} + +// CHECK-LABEL: struct.def @Parent +// CHECK-LABEL: function.def @constrain +// CHECK: %[[AB:.*]] = felt.mul %{{.*}}, %{{.*}} : !felt.type, !felt.type +// CHECK: %[[CHILD:.*]] = struct.readm %{{.*}}[@child] : <@Parent>, !struct.type<@Child> +// CHECK: %[[CALL_AUX:.*]] = struct.readm %{{.*}}[@__llzk_poly_lowering_pass_aux_member_{{[0-9]+}}] : <@Parent>, !felt.type +// CHECK: constrain.eq %[[CALL_AUX]], %[[AB]] : !felt.type, !felt.type +// CHECK: function.call @Child::@constrain(%[[CHILD]], %[[CALL_AUX]]) : (!struct.type<@Child>, !felt.type) -> () From 0ff2d4a6b89d12e9efdafb719c0bde70a5776be8 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:22:49 +0800 Subject: [PATCH 03/12] Topologically order poly auxiliary writes in compute Fixes project-llzk/llzk-lib#551. --- .../unreleased/poly-aux-write-order.yaml | 2 + lib/Transforms/LLZKPolyLoweringPass.cpp | 123 +++++++++++++++++- .../poly_lowering_aux_write_order.llzk | 37 ++++++ 3 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 changelogs/unreleased/poly-aux-write-order.yaml create mode 100644 test/Transforms/PolyLowering/poly_lowering_aux_write_order.llzk diff --git a/changelogs/unreleased/poly-aux-write-order.yaml b/changelogs/unreleased/poly-aux-write-order.yaml new file mode 100644 index 000000000..c8a2ee061 --- /dev/null +++ b/changelogs/unreleased/poly-aux-write-order.yaml @@ -0,0 +1,2 @@ +fixed: + - Topologically order polynomial-lowering auxiliary writes in compute functions diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index 02779e9a9..a9e8c1a21 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -24,7 +24,10 @@ #include #include +#include #include +#include +#include #include #include @@ -51,6 +54,7 @@ namespace { struct AuxAssignment { std::string auxMemberName; Value computedValue; + Value auxValue; }; class PassImpl : public llzk::impl::PolyLoweringPassBase { @@ -82,6 +86,106 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { return defOp->getBlock() == useOp->getBlock() && defOp->isBeforeInBlock(useOp); } + void addAuxDependency( + unsigned dep, unsigned owner, DenseSet &seenDeps, SmallVectorImpl &deps + ) const { + if (dep == owner) { + return; + } + if (seenDeps.insert(dep).second) { + deps.push_back(dep); + } + } + + void collectAuxDependencies( + Value val, unsigned owner, const DenseMap &auxValueToIndex, + const llvm::StringMap &auxNameToIndex, DenseSet &visitedValues, + DenseSet &seenDeps, SmallVectorImpl &deps + ) const { + if (!val || !visitedValues.insert(val).second) { + return; + } + + if (auto it = auxValueToIndex.find(val); it != auxValueToIndex.end()) { + addAuxDependency(it->second, owner, seenDeps, deps); + } + + if (auto readOp = val.getDefiningOp()) { + auto it = auxNameToIndex.find(readOp.getMemberName()); + if (it != auxNameToIndex.end()) { + addAuxDependency(it->second, owner, seenDeps, deps); + } + } + + Operation *defOp = val.getDefiningOp(); + if (!defOp) { + return; + } + + for (Value operand : defOp->getOperands()) { + collectAuxDependencies( + operand, owner, auxValueToIndex, auxNameToIndex, visitedValues, seenDeps, deps + ); + } + } + + LogicalResult visitAuxAssignment( + unsigned idx, ArrayRef> deps, SmallVectorImpl &visitState, + SmallVectorImpl &ordered, ArrayRef auxAssignments + ) const { + if (visitState[idx] == 2) { + return success(); + } + if (visitState[idx] == 1) { + return emitError(auxAssignments[idx].computedValue.getLoc()) + << "poly lowering generated cyclic auxiliary dependency involving @" + << auxAssignments[idx].auxMemberName; + } + + visitState[idx] = 1; + for (unsigned dep : deps[idx]) { + if (failed(visitAuxAssignment(dep, deps, visitState, ordered, auxAssignments))) { + return failure(); + } + } + visitState[idx] = 2; + ordered.push_back(idx); + return success(); + } + + LogicalResult orderAuxAssignments( + ArrayRef auxAssignments, SmallVectorImpl &ordered + ) const { + DenseMap auxValueToIndex; + llvm::StringMap auxNameToIndex; + auxValueToIndex.reserve(auxAssignments.size()); + auxNameToIndex.reserve(auxAssignments.size()); + for (auto [idx, assign] : llvm::enumerate(auxAssignments)) { + if (assign.auxValue) { + auxValueToIndex[assign.auxValue] = idx; + } + auxNameToIndex[assign.auxMemberName] = idx; + } + + SmallVector> deps(auxAssignments.size()); + for (auto [idx, assign] : llvm::enumerate(auxAssignments)) { + DenseSet visitedValues; + DenseSet seenDeps; + collectAuxDependencies( + assign.computedValue, idx, auxValueToIndex, auxNameToIndex, visitedValues, seenDeps, + deps[idx] + ); + } + + SmallVector visitState(auxAssignments.size(), 0); + for (unsigned idx = 0, e = auxAssignments.size(); idx < e; ++idx) { + if (failed(visitAuxAssignment(idx, deps, visitState, ordered, auxAssignments))) { + return failure(); + } + } + return success(); + } + // Recursively compute degree of FeltOps SSA values unsigned getDegree(Value val, DenseMap &memo) { if (auto it = memo.find(val); it != memo.end()) { @@ -214,7 +318,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { auto auxVal = builder.create( lhs.getLoc(), lhs.getType(), selfVal, auxMember.getNameAttr() ); - auxAssignments.push_back({auxName, lhs}); + auxAssignments.push_back({auxName, lhs, auxVal}); Location loc = builder.getFusedLoc({auxVal.getLoc(), lhs.getLoc()}); auto eqOp = builder.create(loc, auxVal, lhs); @@ -247,7 +351,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { // Emit constraint: auxVal == toFactor Location loc = builder.getFusedLoc({auxVal.getLoc(), toFactor.getLoc()}); auto eqOp = builder.create(loc, auxVal, toFactor); - auxAssignments.push_back({auxName, toFactor}); + auxAssignments.push_back({auxName, toFactor, auxVal}); // Update memoization rewrites[toFactor] = auxVal; degreeMemo[auxVal] = 1; // stays same @@ -309,7 +413,7 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { Location loc = builder.getFusedLoc({auxVal.getLoc(), loweredVal.getLoc()}); builder.create(loc, auxVal, loweredVal); - auxAssignments.push_back({auxName, loweredVal}); + auxAssignments.push_back({auxName, loweredVal, auxVal}); degreeMemo[auxVal] = 1; rewrites[loweredVal] = auxVal; @@ -496,13 +600,24 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { OpBuilder builder(&computeBlock, computeBlock.getTerminator()->getIterator()); Value selfVal = computeFunc.getSelfValueFromCompute(); - for (const auto &assign : auxAssignments) { + SmallVector orderedAuxAssignments; + orderedAuxAssignments.reserve(auxAssignments.size()); + if (failed(orderAuxAssignments(auxAssignments, orderedAuxAssignments))) { + signalPassFailure(); + return; + } + + for (unsigned assignIdx : orderedAuxAssignments) { + const auto &assign = auxAssignments[assignIdx]; Value rebuiltExpr = rebuildExprInCompute(assign.computedValue, computeFunc, builder, rebuildMemo); builder.create( assign.computedValue.getLoc(), selfVal, builder.getStringAttr(assign.auxMemberName), rebuiltExpr ); + if (assign.auxValue) { + rebuildMemo[assign.auxValue] = rebuiltExpr; + } } }); } diff --git a/test/Transforms/PolyLowering/poly_lowering_aux_write_order.llzk b/test/Transforms/PolyLowering/poly_lowering_aux_write_order.llzk new file mode 100644 index 000000000..96d8e1eb5 --- /dev/null +++ b/test/Transforms/PolyLowering/poly_lowering_aux_write_order.llzk @@ -0,0 +1,37 @@ +// RUN: llzk-opt -I %S -split-input-file -llzk-full-poly-lowering --verify-each %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang, llzk.main = !struct.type<@CmpConstraint>} { + struct.def @CmpConstraint { + struct.member @out : !felt.type {llzk.pub} + + function.def @compute( + %a: !felt.type {function.arg_name = "a"}, + %b: !felt.type {function.arg_name = "b"} + ) -> !struct.type<@CmpConstraint> { + %self = struct.new : !struct.type<@CmpConstraint> + struct.writem %self[@out] = %a : !struct.type<@CmpConstraint>, !felt.type + function.return %self : !struct.type<@CmpConstraint> + } + + function.def @constrain(%self: !struct.type<@CmpConstraint>, %a: !felt.type, %b: !felt.type) { + %z = felt.mul %a, %b + %zz = felt.mul %z, %z + %za = felt.mul %z, %a + %za_sq = felt.mul %za, %za + constrain.eq %za_sq, %z : !felt.type + constrain.eq %zz, %za : !felt.type + function.return + } + } +} + +// CHECK-LABEL: struct.def @CmpConstraint +// CHECK-LABEL: function.def @compute +// CHECK: %[[SELF:.*]] = struct.new : <@CmpConstraint> +// CHECK-NOT: struct.readm %[[SELF]][@__llzk_poly_lowering_pass_aux_member_ +// CHECK: %[[Z:.*]] = felt.mul %{{.*}}, %{{.*}} : !felt.type, !felt.type +// CHECK: struct.writem %[[SELF]][@__llzk_poly_lowering_pass_aux_member_{{[0-9]+}}] = %[[Z]] : <@CmpConstraint>, !felt.type +// CHECK: %[[ZA:.*]] = felt.mul %[[Z]], %{{.*}} : !felt.type, !felt.type +// CHECK: struct.writem %[[SELF]][@__llzk_poly_lowering_pass_aux_member_{{[0-9]+}}] = %[[ZA]] : <@CmpConstraint>, !felt.type +// CHECK-NOT: struct.readm %[[SELF]][@__llzk_poly_lowering_pass_aux_member_ +// CHECK: function.return %[[SELF]] : !struct.type<@CmpConstraint> From 3300a3248be073171091deb0a8bbe142d5cdb4dd Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:23:20 +0800 Subject: [PATCH 04/12] Preserve member-read access when rebuilding aux producers Fixes project-llzk/llzk-lib#550. --- .../unreleased/poly-member-read-access.yaml | 2 + lib/Transforms/LLZKLoweringUtils.cpp | 61 +++++++++++++++-- .../poly_lowering_member_read_access.llzk | 66 +++++++++++++++++++ .../PolyLowering/poly_lowering_pass_deg2.llzk | 15 +++-- 4 files changed, 130 insertions(+), 14 deletions(-) create mode 100644 changelogs/unreleased/poly-member-read-access.yaml create mode 100644 test/Transforms/PolyLowering/poly_lowering_member_read_access.llzk diff --git a/changelogs/unreleased/poly-member-read-access.yaml b/changelogs/unreleased/poly-member-read-access.yaml new file mode 100644 index 000000000..4c4e11c7b --- /dev/null +++ b/changelogs/unreleased/poly-member-read-access.yaml @@ -0,0 +1,2 @@ +fixed: + - Preserve member-read component operands, access attributes, and index map operands when polynomial lowering rebuilds auxiliary producers diff --git a/lib/Transforms/LLZKLoweringUtils.cpp b/lib/Transforms/LLZKLoweringUtils.cpp index 15401eff6..209a778cc 100644 --- a/lib/Transforms/LLZKLoweringUtils.cpp +++ b/lib/Transforms/LLZKLoweringUtils.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -26,6 +27,33 @@ using namespace llzk::constrain; namespace llzk { +namespace { + +Value mapBlockArgumentInCompute(BlockArgument barg, FuncDefOp computeFunc) { + // constrain(%self, args...) corresponds to compute(args...). When rebuilding + // expressions in compute(), the constrain self argument maps to the compute + // self value, and every other constrain argument maps one position earlier in + // compute(). + if (barg.getArgNumber() == 0) { + return computeFunc.getSelfValueFromCompute(); + } + return computeFunc.getArgument(barg.getArgNumber() - 1); +} + +Value mapValueIntoCompute( + Value val, FuncDefOp computeFunc, OpBuilder &builder, DenseMap &memo +) { + if (auto it = memo.find(val); it != memo.end()) { + return it->second; + } + if (auto barg = llvm::dyn_cast(val)) { + return memo[val] = mapBlockArgumentInCompute(barg, computeFunc); + } + return rebuildExprInCompute(val, computeFunc, builder, memo); +} + +} // namespace + Value rebuildExprInCompute( Value val, FuncDefOp computeFunc, OpBuilder &builder, DenseMap &memo ) { @@ -34,17 +62,36 @@ Value rebuildExprInCompute( } if (auto barg = llvm::dyn_cast(val)) { - unsigned index = barg.getArgNumber(); - Value mapped = computeFunc.getArgument(index - 1); - return memo[val] = mapped; + return memo[val] = mapBlockArgumentInCompute(barg, computeFunc); } if (auto readOp = val.getDefiningOp()) { - Value self = computeFunc.getSelfValueFromCompute(); - Value rebuilt = builder.create( - readOp.getLoc(), readOp.getType(), self, readOp.getMemberNameAttr().getAttr() + IRMapping mapper; + for (Value operand : readOp->getOperands()) { + mapper.map(operand, mapValueIntoCompute(operand, computeFunc, builder, memo)); + } + + Operation *rebuiltOp = builder.clone(*readOp.getOperation(), mapper); + assert(rebuiltOp->getNumResults() == 1 && "member reads have exactly one result"); + return memo[val] = rebuiltOp->getResult(0); + } + + if (val.getType().isIndex()) { + Operation *defOp = val.getDefiningOp(); + assert(defOp && "index block arguments should already be mapped"); + + IRMapping mapper; + for (Value operand : defOp->getOperands()) { + mapper.map(operand, mapValueIntoCompute(operand, computeFunc, builder, memo)); + } + + Operation *rebuiltOp = builder.clone(*defOp, mapper); + assert( + rebuiltOp->getNumResults() == defOp->getNumResults() && + "cloned index op should preserve result count" ); - return memo[val] = rebuilt; + unsigned resultNumber = llvm::cast(val).getResultNumber(); + return memo[val] = rebuiltOp->getResult(resultNumber); } if (auto add = val.getDefiningOp()) { diff --git a/test/Transforms/PolyLowering/poly_lowering_member_read_access.llzk b/test/Transforms/PolyLowering/poly_lowering_member_read_access.llzk new file mode 100644 index 000000000..47ec0bfb6 --- /dev/null +++ b/test/Transforms/PolyLowering/poly_lowering_member_read_access.llzk @@ -0,0 +1,66 @@ +// RUN: llzk-opt -I %S -split-input-file -llzk-poly-lowering-pass="max-degree=2" --verify-each %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + struct.def @TableOffsetAux { + struct.member @x : !felt.type {column} + struct.member @out : !felt.type {llzk.pub} + + function.def @compute(%x: !felt.type) -> !struct.type<@TableOffsetAux> { + %self = struct.new : !struct.type<@TableOffsetAux> + struct.writem %self[@x] = %x : !struct.type<@TableOffsetAux>, !felt.type + function.return %self : !struct.type<@TableOffsetAux> + } + + function.def @constrain(%self: !struct.type<@TableOffsetAux>, %x: !felt.type) { + %prev = struct.readm %self[@x] : !struct.type<@TableOffsetAux>, !felt.type {tableOffset = -1 : index} + %sq = felt.mul %prev, %prev : !felt.type, !felt.type + %fourth = felt.mul %sq, %sq : !felt.type, !felt.type + %out = struct.readm %self[@out] : !struct.type<@TableOffsetAux>, !felt.type + constrain.eq %out, %fourth : !felt.type + function.return + } + } +} + +// CHECK-LABEL: struct.def @TableOffsetAux +// CHECK-LABEL: function.def @compute +// CHECK: %[[SELF:.*]] = struct.new : <@TableOffsetAux> +// CHECK-NOT: struct.readm %[[SELF]][@x] : <@TableOffsetAux>, !felt.type{{$}} +// CHECK: struct.readm %[[SELF]][@x] : <@TableOffsetAux>, !felt.type {tableOffset = -1 : index} +// CHECK: struct.writem %[[SELF]][@__llzk_poly_lowering_pass_aux_member_0] +// CHECK-LABEL: function.def @constrain +// CHECK: struct.readm {{.*}}[@x] : <@TableOffsetAux>, !felt.type {tableOffset = -1 : index} +// CHECK: struct.readm {{.*}}[@__llzk_poly_lowering_pass_aux_member_0] : <@TableOffsetAux>, !felt.type +// ----- + +module attributes {llzk.lang} { + struct.def @AffineOffsetAux { + struct.member @x : !felt.type {column} + struct.member @out : !felt.type {llzk.pub} + + function.def @compute(%x: !felt.type) -> !struct.type<@AffineOffsetAux> { + %self = struct.new : !struct.type<@AffineOffsetAux> + struct.writem %self[@x] = %x : !struct.type<@AffineOffsetAux>, !felt.type + function.return %self : !struct.type<@AffineOffsetAux> + } + + function.def @constrain(%self: !struct.type<@AffineOffsetAux>, %x: !felt.type) { + %idx = arith.constant 2 : index + %prev = struct.readm %self[@x] {()[%idx]} : !struct.type<@AffineOffsetAux>, !felt.type {tableOffset = affine_map<()[s0] -> (s0 - 1)>} + %sq = felt.mul %prev, %prev : !felt.type, !felt.type + %fourth = felt.mul %sq, %sq : !felt.type, !felt.type + %out = struct.readm %self[@out] : !struct.type<@AffineOffsetAux>, !felt.type + constrain.eq %out, %fourth : !felt.type + function.return + } + } +} + +// CHECK-LABEL: struct.def @AffineOffsetAux +// CHECK-LABEL: function.def @compute +// CHECK: %[[SELF2:.*]] = struct.new : <@AffineOffsetAux> +// CHECK: %[[IDX:.*]] = arith.constant 2 : index +// CHECK: struct.readm %[[SELF2]][@x] {()[%[[IDX]]]} : <@AffineOffsetAux>, !felt.type {tableOffset = {{.*}}} +// CHECK: struct.writem %[[SELF2]][@__llzk_poly_lowering_pass_aux_member_{{[0-9]+}}] +// CHECK-LABEL: function.def @constrain +// CHECK: struct.readm {{.*}}[@x] {()[%{{.*}}]} : <@AffineOffsetAux>, !felt.type {tableOffset = {{.*}}} diff --git a/test/Transforms/PolyLowering/poly_lowering_pass_deg2.llzk b/test/Transforms/PolyLowering/poly_lowering_pass_deg2.llzk index 0e3a18a72..215125daf 100644 --- a/test/Transforms/PolyLowering/poly_lowering_pass_deg2.llzk +++ b/test/Transforms/PolyLowering/poly_lowering_pass_deg2.llzk @@ -203,14 +203,15 @@ module attributes {llzk.lang} { // CHECK: struct.member @val : !felt.type // CHECK: function.def @compute(%[[VAL_0:.*]]: !felt.type, %[[VAL_1:.*]]: !felt.type) -> !struct.type<@Mod1> attributes {function.allow_witness} { // CHECK: %[[VAL_2:.*]] = struct.new : <@Mod1> -// CHECK: %[[VAL_3:.*]] = struct.readm %[[VAL_2]][@val] : <@Mod1>, !felt.type -// CHECK: %[[VAL_4:.*]] = felt.mul %[[VAL_3]], %[[VAL_3]] : !felt.type, !felt.type -// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_1] = %[[VAL_4]] : <@Mod1>, !felt.type -// CHECK: %[[VAL_5:.*]] = felt.mul %[[VAL_0]], %[[VAL_1]] : !felt.type, !felt.type -// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_2] = %[[VAL_5]] : <@Mod1>, !felt.type -// CHECK: %[[VAL_7:.*]] = felt.mul %[[VAL_5]], %[[VAL_5]] : !felt.type, !felt.type +// CHECK: %[[VAL_3:.*]] = struct.readm %[[VAL_2]][@mod2] : <@Mod1>, !struct.type<@Mod2> +// CHECK: %[[VAL_4:.*]] = struct.readm %[[VAL_3]][@val] : <@Mod2>, !felt.type +// CHECK: %[[VAL_5:.*]] = felt.mul %[[VAL_4]], %[[VAL_4]] : !felt.type, !felt.type +// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_1] = %[[VAL_5]] : <@Mod1>, !felt.type +// CHECK: %[[VAL_6:.*]] = felt.mul %[[VAL_0]], %[[VAL_1]] : !felt.type, !felt.type +// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_2] = %[[VAL_6]] : <@Mod1>, !felt.type +// CHECK: %[[VAL_7:.*]] = felt.mul %[[VAL_6]], %[[VAL_6]] : !felt.type, !felt.type // CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_3] = %[[VAL_7]] : <@Mod1>, !felt.type -// CHECK: %[[VAL_8:.*]] = felt.mul %[[VAL_5]], %[[VAL_0]] : !felt.type, !felt.type +// CHECK: %[[VAL_8:.*]] = felt.mul %[[VAL_6]], %[[VAL_0]] : !felt.type, !felt.type // CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_4] = %[[VAL_8]] : <@Mod1>, !felt.type // CHECK: function.return %[[VAL_2]] : !struct.type<@Mod1> // CHECK: } From c5c92aecb859d1f9c3585cc0a51b272e10134e64 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:25:17 +0800 Subject: [PATCH 05/12] Reject non-straight-line poly compute bodies Fixes project-llzk/llzk-lib#554. --- .../poly-compute-straight-line.yaml | 2 ++ include/llzk/Transforms/LLZKLoweringUtils.h | 14 +++++++++---- .../Transforms/LLZKTransformationPasses.td | 8 +++---- lib/Transforms/LLZKLoweringUtils.cpp | 18 +++++++++------- lib/Transforms/LLZKPolyLoweringPass.cpp | 5 +++++ ...poly_lowering_fail_compute_multiblock.llzk | 21 +++++++++++++++++++ 6 files changed, 53 insertions(+), 15 deletions(-) create mode 100644 changelogs/unreleased/poly-compute-straight-line.yaml create mode 100644 test/Transforms/PolyLowering/poly_lowering_fail_compute_multiblock.llzk diff --git a/changelogs/unreleased/poly-compute-straight-line.yaml b/changelogs/unreleased/poly-compute-straight-line.yaml new file mode 100644 index 000000000..6fc046379 --- /dev/null +++ b/changelogs/unreleased/poly-compute-straight-line.yaml @@ -0,0 +1,2 @@ +fixed: + - Reject non-straight-line compute bodies before polynomial lowering inserts auxiliary writes diff --git a/include/llzk/Transforms/LLZKLoweringUtils.h b/include/llzk/Transforms/LLZKLoweringUtils.h index 8ccaee945..73dc4df14 100644 --- a/include/llzk/Transforms/LLZKLoweringUtils.h +++ b/include/llzk/Transforms/LLZKLoweringUtils.h @@ -33,10 +33,16 @@ mlir::Value rebuildExprInCompute( mlir::LogicalResult checkForAuxMemberConflicts(component::StructDefOp structDef, llvm::StringRef auxPrefix); -/// Rejects control flow under `constrainFunc`; polynomial and R1CS auxiliary -/// materialization assumes control flow has already been flattened or otherwise -/// lowered away. The region check catches multi-block function bodies before -/// the operation walk rejects nested regions or successor-bearing operations. +/// Rejects control flow under `func`; auxiliary materialization assumes control +/// flow has already been flattened or otherwise lowered away. The region check +/// catches multi-block function bodies before the operation walk rejects nested +/// regions or successor-bearing operations. +mlir::LogicalResult checkFuncBodyIsStraightLine( + function::FuncDefOp func, llvm::StringRef passName, llvm::StringRef funcName +); + +/// Rejects control flow under `constrainFunc`; kept as a convenience wrapper +/// for passes whose public diagnostic already names constrain bodies. mlir::LogicalResult checkConstrainBodyIsStraightLine(function::FuncDefOp constrainFunc, llvm::StringRef passName); diff --git a/include/llzk/Transforms/LLZKTransformationPasses.td b/include/llzk/Transforms/LLZKTransformationPasses.td index 8172f80d4..c1493fa23 100644 --- a/include/llzk/Transforms/LLZKTransformationPasses.td +++ b/include/llzk/Transforms/LLZKTransformationPasses.td @@ -74,10 +74,10 @@ def PolyLoweringPass : LLZKPass<"llzk-poly-lowering-pass"> { through degree-neutral roots such as `felt.add`, `felt.sub`, and `felt.neg`, so composite equality operands and struct constrain call arguments satisfy their degree bounds. - The pass expects already-flattened, straight-line `constrain` bodies. Run `llzk-flatten` or - another control-flow lowering pass before this pass; nested-region, successor-bearing, or - multi-block constraints are rejected instead of being rewritten with component-lifetime - auxiliary members. + The pass expects already-flattened, straight-line `compute` and `constrain` bodies. Run + `llzk-flatten` or another control-flow lowering pass before this pass; nested-region, + successor-bearing, or multi-block functions are rejected instead of being rewritten with + component-lifetime auxiliary members. This pass is best used on already-flattened input as part of the `-llzk-full-poly-lowering` pipeline, which includes additional cleanup passes to ensure correctness and optimal diff --git a/lib/Transforms/LLZKLoweringUtils.cpp b/lib/Transforms/LLZKLoweringUtils.cpp index 209a778cc..e94ca8385 100644 --- a/lib/Transforms/LLZKLoweringUtils.cpp +++ b/lib/Transforms/LLZKLoweringUtils.cpp @@ -146,16 +146,16 @@ LogicalResult checkForAuxMemberConflicts(StructDefOp structDef, StringRef prefix return failure(conflictFound); } -LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRef passName) { - auto emitStraightLineError = [passName](Operation *op) { - op->emitError() << passName - << " expects a straight-line constrain body; run `llzk-flatten` or another " - "control-flow lowering pass first"; +LogicalResult +checkFuncBodyIsStraightLine(FuncDefOp func, StringRef passName, StringRef funcName) { + auto emitStraightLineError = [passName, funcName](Operation *op) { + op->emitError() << passName << " expects a straight-line " << funcName + << " body; run `llzk-flatten` or another control-flow lowering pass first"; }; - Region &body = constrainFunc.getBody(); + Region &body = func.getBody(); if (!body.hasOneBlock()) { - emitStraightLineError(constrainFunc.getOperation()); + emitStraightLineError(func.getOperation()); return failure(); } @@ -176,6 +176,10 @@ LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRe return failure(); } +LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRef passName) { + return checkFuncBodyIsStraightLine(constrainFunc, passName, "constrain"); +} + void replaceSubsequentUsesWith(Value oldVal, Value newVal, Operation *afterOp) { assert(afterOp && "afterOp must be a valid Operation*"); diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index a9e8c1a21..e09aada55 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -514,6 +514,11 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { return; } + if (failed(checkFuncBodyIsStraightLine(computeFunc, "poly lowering", "compute"))) { + signalPassFailure(); + return; + } + DenseMap degreeMemo; DenseMap rewrites; SmallVector auxAssignments; diff --git a/test/Transforms/PolyLowering/poly_lowering_fail_compute_multiblock.llzk b/test/Transforms/PolyLowering/poly_lowering_fail_compute_multiblock.llzk new file mode 100644 index 000000000..20b5b1653 --- /dev/null +++ b/test/Transforms/PolyLowering/poly_lowering_fail_compute_multiblock.llzk @@ -0,0 +1,21 @@ +// RUN: llzk-opt -I %S -split-input-file -llzk-poly-lowering-pass="max-degree=2" -verify-diagnostics %s + +module attributes {llzk.lang = "llzk"} { + struct.def @ComputeMultiBlockAuxScope { + // expected-error@+1 {{poly lowering expects a straight-line compute body; run `llzk-flatten` or another control-flow lowering pass first}} + function.def @compute(%a: !felt.type, %b: !felt.type) -> !struct.type<@ComputeMultiBlockAuxScope> { + %entry_self = struct.new : !struct.type<@ComputeMultiBlockAuxScope> + function.return %entry_self : !struct.type<@ComputeMultiBlockAuxScope> + ^late: + %late_self = struct.new : !struct.type<@ComputeMultiBlockAuxScope> + function.return %late_self : !struct.type<@ComputeMultiBlockAuxScope> + } + + function.def @constrain(%self: !struct.type<@ComputeMultiBlockAuxScope>, %a: !felt.type, %b: !felt.type) { + %ab = felt.mul %a, %b : !felt.type, !felt.type + %cube = felt.mul %ab, %a : !felt.type, !felt.type + constrain.eq %cube, %a : !felt.type + function.return + } + } +} From dadff5ddfb3e9b2f9bfae67b87022a49e0fd8198 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:35:25 +0800 Subject: [PATCH 06/12] Format poly lowering fixes for CI --- lib/Transforms/LLZKLoweringUtils.cpp | 3 +-- lib/Transforms/LLZKPolyLoweringPass.cpp | 10 ++++------ 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/lib/Transforms/LLZKLoweringUtils.cpp b/lib/Transforms/LLZKLoweringUtils.cpp index e94ca8385..d7cd3d199 100644 --- a/lib/Transforms/LLZKLoweringUtils.cpp +++ b/lib/Transforms/LLZKLoweringUtils.cpp @@ -146,8 +146,7 @@ LogicalResult checkForAuxMemberConflicts(StructDefOp structDef, StringRef prefix return failure(conflictFound); } -LogicalResult -checkFuncBodyIsStraightLine(FuncDefOp func, StringRef passName, StringRef funcName) { +LogicalResult checkFuncBodyIsStraightLine(FuncDefOp func, StringRef passName, StringRef funcName) { auto emitStraightLineError = [passName, funcName](Operation *op) { op->emitError() << passName << " expects a straight-line " << funcName << " body; run `llzk-flatten` or another control-flow lowering pass first"; diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index e09aada55..e45bcd963 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -25,8 +25,8 @@ #include #include #include -#include #include +#include #include #include @@ -390,11 +390,9 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { DenseMap °reeMemo, DenseMap &rewrites, SmallVector &auxAssignments ) { - Value loweredVal = - lowerExpression( - val, structDef, constrainFunc, callOp.getOperation(), degreeMemo, rewrites, - auxAssignments - ); + Value loweredVal = lowerExpression( + val, structDef, constrainFunc, callOp.getOperation(), degreeMemo, rewrites, auxAssignments + ); DenseMap checkMemo; if (getDegree(loweredVal, checkMemo) <= 1) { return loweredVal; From 903431dc6941e44bb1b307ff8170fa3e1e6425cc Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:43:37 +0800 Subject: [PATCH 07/12] Remove unsupported StringMap preallocation --- lib/Transforms/LLZKPolyLoweringPass.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/Transforms/LLZKPolyLoweringPass.cpp b/lib/Transforms/LLZKPolyLoweringPass.cpp index e45bcd963..131f8f447 100644 --- a/lib/Transforms/LLZKPolyLoweringPass.cpp +++ b/lib/Transforms/LLZKPolyLoweringPass.cpp @@ -159,7 +159,6 @@ class PassImpl : public llzk::impl::PolyLoweringPassBase { DenseMap auxValueToIndex; llvm::StringMap auxNameToIndex; auxValueToIndex.reserve(auxAssignments.size()); - auxNameToIndex.reserve(auxAssignments.size()); for (auto [idx, assign] : llvm::enumerate(auxAssignments)) { if (assign.auxValue) { auxValueToIndex[assign.auxValue] = idx; From 05c057eac5f17aef3c180d86a2fa824df955b376 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:56:14 +0800 Subject: [PATCH 08/12] Update poly lowering regression expectations --- .../PolyLowering/poly_lowering_composite_roots.llzk | 3 +-- test/Transforms/PolyLowering/poly_lowering_pass_deg3.llzk | 5 ++--- .../R1CSLowering/r1cs_lowering_typed_aux_member.llzk | 2 ++ 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/test/Transforms/PolyLowering/poly_lowering_composite_roots.llzk b/test/Transforms/PolyLowering/poly_lowering_composite_roots.llzk index fb2aae76d..7189339aa 100644 --- a/test/Transforms/PolyLowering/poly_lowering_composite_roots.llzk +++ b/test/Transforms/PolyLowering/poly_lowering_composite_roots.llzk @@ -167,8 +167,7 @@ module attributes {llzk.lang} { // CHECK: %[[C_SELF:.*]] = struct.new : <@CompositeCallArg> // CHECK: %[[C_AB:.*]] = felt.mul %[[C_A]], %[[C_B]] : !felt.type, !felt.type // CHECK: struct.writem %[[C_SELF]][@__llzk_poly_lowering_pass_aux_member_0] = %[[C_AB]] : <@CompositeCallArg>, !felt.type -// CHECK: %[[C_AUX0:.*]] = struct.readm %[[C_SELF]][@__llzk_poly_lowering_pass_aux_member_0] : <@CompositeCallArg>, !felt.type -// CHECK: %[[C_ABC:.*]] = felt.mul %[[C_AUX0]], %[[C_C]] : !felt.type, !felt.type +// CHECK: %[[C_ABC:.*]] = felt.mul %[[C_AB]], %[[C_C]] : !felt.type, !felt.type // CHECK: %[[C_ARG:.*]] = felt.add %[[C_ABC]], %[[C_D]] : !felt.type, !felt.type // CHECK: struct.writem %[[C_SELF]][@__llzk_poly_lowering_pass_aux_member_1] = %[[C_ARG]] : <@CompositeCallArg>, !felt.type // CHECK: function.return %[[C_SELF]] : !struct.type<@CompositeCallArg> diff --git a/test/Transforms/PolyLowering/poly_lowering_pass_deg3.llzk b/test/Transforms/PolyLowering/poly_lowering_pass_deg3.llzk index d43c2e124..afe8eff7e 100644 --- a/test/Transforms/PolyLowering/poly_lowering_pass_deg3.llzk +++ b/test/Transforms/PolyLowering/poly_lowering_pass_deg3.llzk @@ -23,11 +23,10 @@ module attributes {llzk.lang} { // CHECK-LABEL: struct.def @CmpConstraint { // CHECK: function.def @compute(%[[VAL_0:.*]]: !felt.type, %[[VAL_1:.*]]: !felt.type) -> !struct.type<@CmpConstraint> attributes {function.allow_witness} { // CHECK: %[[VAL_2:.*]] = struct.new : <@CmpConstraint> -// CHECK: %[[VAL_3:.*]] = struct.readm %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_1] : <@CmpConstraint>, !felt.type +// CHECK: %[[VAL_3:.*]] = felt.mul %[[VAL_0]], %[[VAL_1]] : !felt.type, !felt.type +// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_1] = %[[VAL_3]] : <@CmpConstraint>, !felt.type // CHECK: %[[VAL_4:.*]] = felt.mul %[[VAL_3]], %[[VAL_0]] : !felt.type, !felt.type // CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_0] = %[[VAL_4]] : <@CmpConstraint>, !felt.type -// CHECK: %[[VAL_5:.*]] = felt.mul %[[VAL_0]], %[[VAL_1]] : !felt.type, !felt.type -// CHECK: struct.writem %[[VAL_2]][@__llzk_poly_lowering_pass_aux_member_1] = %[[VAL_5]] : <@CmpConstraint>, !felt.type // CHECK: function.return %[[VAL_2]] : !struct.type<@CmpConstraint> // CHECK: } // CHECK: function.def @constrain(%[[VAL_6:.*]]: !struct.type<@CmpConstraint>, %[[VAL_7:.*]]: !felt.type, %[[VAL_8:.*]]: !felt.type) attributes {function.allow_constraint} { diff --git a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk index fc62e3f09..82c70a514 100644 --- a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk +++ b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk @@ -22,6 +22,8 @@ module attributes {llzk.lang} { %ab = felt.mul %a, %b : !felt.type<"babybear">, !felt.type<"babybear"> %bc = felt.mul %b, %c : !felt.type<"babybear">, !felt.type<"babybear"> constrain.eq %ab, %bc : !felt.type<"babybear"> + %out = struct.readm %self[@out] : !struct.type<@TypedR1CSAux>, !felt.type<"babybear"> + constrain.eq %out, %ab : !felt.type<"babybear"> function.return } } From 0d6bdf6cc5e90d94dec125f3fb0eccae5e53659b Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:08:15 +0800 Subject: [PATCH 09/12] Write typed R1CS public output in regression --- .../Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk | 2 ++ 1 file changed, 2 insertions(+) diff --git a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk index 82c70a514..c413710f5 100644 --- a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk +++ b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk @@ -10,6 +10,8 @@ module attributes {llzk.lang} { %c: !felt.type<"babybear"> ) -> !struct.type<@TypedR1CSAux> { %self = struct.new : !struct.type<@TypedR1CSAux> + %ab = felt.mul %a, %b : !felt.type<"babybear">, !felt.type<"babybear"> + struct.writem %self[@out] = %ab : !struct.type<@TypedR1CSAux>, !felt.type<"babybear"> function.return %self : !struct.type<@TypedR1CSAux> } From 09c57ecbbf95de0b4274ae0b819e6b02cce49f4b Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:24:29 +0800 Subject: [PATCH 10/12] Use typed zero constants in R1CS lowering --- backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp b/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp index cbd88787a..880b02f39 100644 --- a/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp +++ b/backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp @@ -532,7 +532,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase { // Entire linear combination was zero result = builder.create( loc, r1cs::LinearType::get(builder.getContext()), - r1cs::FeltAttr::get(builder.getContext(), 0) + r1cs::FeltAttr::get(builder.getContext(), toAPSInt(lc.constant)) ); } From f7e642c4471d3d3cbe60b6c3f5fd2faf90f96140 Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:37:34 +0800 Subject: [PATCH 11/12] Tighten typed R1CS FileCheck pattern --- .../R1CSLowering/r1cs_lowering_typed_aux_member.llzk | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk index c413710f5..3a0e0abae 100644 --- a/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk +++ b/test/Transforms/R1CSLowering/r1cs_lowering_typed_aux_member.llzk @@ -32,7 +32,7 @@ module attributes {llzk.lang} { } // CHECK-LABEL: r1cs.circuit @TypedR1CSAux inputs -// CHECK-SAME: %{{.*}}: !r1cs.signal {#r1cs.pub} -// CHECK-SAME: %{{.*}}: !r1cs.signal -// CHECK-SAME: %{{.*}}: !r1cs.signal +// CHECK-SAME: %{{[0-9a-zA-Z_\.]+}}: !r1cs.signal {#r1cs.pub} +// CHECK-SAME: %{{[0-9a-zA-Z_\.]+}}: !r1cs.signal +// CHECK-SAME: %{{[0-9a-zA-Z_\.]+}}: !r1cs.signal // CHECK: r1cs.constrain From 51d4387562f065f8257285ec23f706af90439a2c Mon Sep 17 00:00:00 2001 From: "Cyne Jarvis J. Zarceno" <165563006+Kuhai9801@users.noreply.github.com> Date: Wed, 17 Jun 2026 20:39:48 +0800 Subject: [PATCH 12/12] Document R1CS zero constant fix in changelog --- changelogs/unreleased/poly-typed-aux-members.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/changelogs/unreleased/poly-typed-aux-members.yaml b/changelogs/unreleased/poly-typed-aux-members.yaml index c73ef93c5..d4561d768 100644 --- a/changelogs/unreleased/poly-typed-aux-members.yaml +++ b/changelogs/unreleased/poly-typed-aux-members.yaml @@ -1,3 +1,4 @@ fixed: - Emit polynomial-lowering auxiliary members with the exact type of the materialized expression - Emit all R1CS-lowering auxiliary members with the exact type of the materialized expression + - Emit synthesized zero R1CS linear-combination constants with a printable integer width