diff --git a/changelogs/unreleased/th__296.yaml b/changelogs/unreleased/th__296.yaml new file mode 100644 index 0000000000..6dcf779cf3 --- /dev/null +++ b/changelogs/unreleased/th__296.yaml @@ -0,0 +1,2 @@ +fixed: + - "Fix issue #296 so flattening pass creates multiple instantiated structs for heterogeneous arrays of structs." diff --git a/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td b/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td index d4dd70eb6e..fc20b14fa0 100644 --- a/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td +++ b/include/llzk/Dialect/Polymorphic/Transforms/TransformationPasses.td @@ -56,12 +56,22 @@ def FlatteningPass : LLZKPass<"llzk-flatten"> { let summary = "Flatten structs and unroll loops"; let description = [{ Performs the following transformations: - - Instantiate `affine_map` parameters of StructType and ArrayType - to constant values using the arguments at the instantiation site - - Replace parameterized structs with flattened (i.e., no parameter) - versions of those structs based on requested return type at calls - to `compute()` functions and unroll loops - - Unroll loops + - Instantiate the main struct, when one is specified. + - Instantiate parameterized structs and free functions when concrete + template arguments are available at use sites. + - Replace instantiated `poly.const_read` uses, symbolic member table + offsets, wildcard array reads/writes, and callees rooted at concrete + template parameters inside cloned bodies. + - Unroll loops with statically-known trip counts. + - Fold `affine_map` parameters of StructType and ArrayType to constant + values using the operands at the instantiation site. + - Scalarize pseudo-homogeneous arrays whose static elements refine to + different concrete types after instantiation. + - Propagate refined types through arrays, members, calls, returns, inferred + result types, and unifiable casts until a fixpoint is reached. + - Clean up unused parameterized definitions according to the selected + cleanup mode, remove empty templates, and remove unused discardable + array allocations. }]; // Implementation note: These options should be kept in sync with // `StructInliningFlatteningOptions` in `LLZKTransformationPassPipelines.h`. diff --git a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp index cf69e557c7..1afd8fc21d 100644 --- a/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp +++ b/lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp @@ -21,6 +21,7 @@ #include "llzk/Dialect/Function/IR/Ops.h" #include "llzk/Dialect/LLZK/IR/AttributeHelper.h" #include "llzk/Dialect/LLZK/IR/Attrs.h" +#include "llzk/Dialect/LLZK/IR/Ops.h" #include "llzk/Dialect/Polymorphic/IR/Ops.h" #include "llzk/Dialect/Polymorphic/Transforms/TransformationPasses.h" #include "llzk/Dialect/String/IR/Dialect.h" @@ -43,6 +44,7 @@ #include #include #include +#include #include #include #include @@ -56,8 +58,10 @@ #include #include #include +#include #include #include +#include #include @@ -83,6 +87,8 @@ using namespace llzk::polymorphic::detail; namespace { +/// Emit diagnostics that were collected while converting a cloned body, rebasing placeholder notes +/// onto the call site that triggered the clone. static void reportDelayedDiagnostics(CallOp caller, SmallVector &&diagnostics) { DiagnosticEngine &engine = caller.getContext()->getDiagEngine(); for (Diagnostic &diag : diagnostics) { @@ -109,8 +115,14 @@ class ConversionTracker { StringAttr functionName; }; + /// Published result of one successful full-function conversion. + struct FullFuncInstantiation { + ArrayAttr concreteParamKey; + SymbolRefAttr functionPath; + }; + /// Tracks if some step performed a modification of the code such that another pass should be run. - bool modified; + bool modified = false; /// Maps original remote (i.e., use site) type to new remote type. /// Note: The keys are always parameterized StructType and the values are no-parameter StructType. DenseMap structInstantiations; @@ -124,12 +136,23 @@ class ConversionTracker { /// Maps new remote type (i.e., the values in 'structInstantiations') to a list of Diagnostic /// to report at the location(s) of the compute() that causes the instantiation to the StructType. DenseMap> delayedDiagnostics; + /// Successful full functions keyed by their source operation and exact concrete bindings. + /// Rendered symbol names are presentation-only and can collide for valid parameter values. + /// This cache must outlive individual rewrite pattern instances because flattening re-runs + /// Step 2 across fixpoint iterations. + DenseMap> fullFuncInstantiations; public: + /// Return whether the current flattening iteration has changed the IR. bool isModified() const { return modified; } + + /// Clear the per-iteration modification flag before starting the next iteration. void resetModifiedFlag() { modified = false; } + + /// Merge the modification status from one rewrite step into the iteration state. void updateModifiedFlag(bool currStepModified) { modified |= currStepModified; } + /// Record a struct instantiation from the original use-site type to its cloned replacement type. void recordInstantiation(StructType oldType, StructType newType) { assert(!isNullOrEmpty(oldType.getParams()) && "cannot instantiate with no params"); @@ -159,6 +182,15 @@ class ConversionTracker { return std::nullopt; } + /// Return the original parameterized type that produced the given instantiated type, if any. + std::optional getPreimage(StructType newType) const { + auto cachedResult = reverseInstantiations.find(newType); + if (cachedResult != reverseInstantiations.end()) { + return cachedResult->second; + } + return std::nullopt; + } + /// Record that the given free function was instantiated. void recordInstantiation(SymbolRefAttr funcName) { funcInstantiations.insert(funcName); @@ -214,6 +246,7 @@ class ConversionTracker { return instantiatedNames; } + /// Emit diagnostics delayed until a compute call has been rewritten to the instantiated type. void reportDelayedDiagnostics(StructType newType, CallOp caller) { auto res = delayedDiagnostics.find(newType); if (res != delayedDiagnostics.end()) { @@ -226,10 +259,39 @@ class ConversionTracker { } } + /// Return the mutable diagnostic queue associated with `newType`. SmallVector &delayedDiagnosticSet(StructType newType) { return delayedDiagnostics[newType]; } + /// Return the successfully converted full function for this exact source/key pair, if any. + std::optional + getFullFuncInstantiation(FuncDefOp sourceFunc, ArrayAttr concreteParamKey) const { + auto found = fullFuncInstantiations.find(sourceFunc.getOperation()); + if (found == fullFuncInstantiations.end()) { + return std::nullopt; + } + for (const FullFuncInstantiation &candidate : found->second) { + if (candidate.concreteParamKey == concreteParamKey) { + return candidate.functionPath; + } + } + return std::nullopt; + } + + /// Publish a successful full conversion after insertion and body conversion have completed. + void recordFullFuncInstantiation( + FuncDefOp sourceFunc, ArrayAttr concreteParamKey, SymbolRefAttr instantiatedPath + ) { + assert( + !getFullFuncInstantiation(sourceFunc, concreteParamKey).has_value() && + "full function instantiation already cached" + ); + fullFuncInstantiations[sourceFunc.getOperation()].push_back( + FullFuncInstantiation {concreteParamKey, instantiatedPath} + ); + } + /// Check if the type conversion is legal, i.e., the new type unifies with and is more concrete /// than the old type with additional allowance for the results of struct flattening conversions. bool isLegalConversion(Type oldType, Type newType, const char *patName) const { @@ -266,6 +328,7 @@ class ConversionTracker { return false; } + /// Check whether every corresponding pair in `oldTypes` and `newTypes` is a legal conversion. template inline bool areLegalConversions(T oldTypes, U newTypes, const char *patName) const { return llvm::all_of( @@ -276,11 +339,14 @@ class ConversionTracker { } }; +/// Base conversion pattern for ops that reference template symbols by attribute and rewrite only +/// when that symbol has a concrete instantiation value of one of `HandledAttrs`. template class SymbolUserHelper : public OpConversionPattern { private: const DenseMap ¶mNameToValue; + /// Construct the CRTP helper with the template binding map used for symbol lookups. SymbolUserHelper( TypeConverter &converter, MLIRContext *ctx, unsigned patternBenefit, const DenseMap ¶mNameToInstantiatedValue @@ -291,14 +357,17 @@ class SymbolUserHelper : public OpConversionPattern { public: using OpAdaptor = typename mlir::OpConversionPattern::OpAdaptor; + /// Return the attribute on `op` that should be looked up in the instantiation map. virtual Attribute getNameAttr(Op) const = 0; + /// Report a type-specific fallback diagnostic for instantiated values not handled by `Impl`. virtual LogicalResult handleDefaultRewrite( Attribute, Op op, OpAdaptor, ConversionPatternRewriter &, Attribute a ) const { return op->emitOpError().append("expected value with type ", op.getType(), " but found ", a); } + /// Dispatch an instantiated symbol value to the concrete `Impl::handleRewrite` overload. LogicalResult matchAndRewrite(Op op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { LLVM_DEBUG(llvm::dbgs() << "[SymbolUserHelper] op: " << op << '\n'); @@ -322,6 +391,35 @@ class SymbolUserHelper : public OpConversionPattern { friend Impl; }; +/// Materialize a template value as a constant suitable for `resultType`, or return a null value +/// when the value cannot be represented by one of the supported constant operations. +/// +/// Callers that need diagnostics for lossy conversions (such as non-zero integers to `i1`) remain +/// responsible for reporting them before calling this helper. +static Value +materializeTemplateConstant(OpBuilder &builder, Location loc, Type resultType, Attribute value) { + if (IntegerAttr integer = llvm::dyn_cast(value)) { + if (FeltType type = llvm::dyn_cast(resultType)) { + return builder.create( + loc, FeltConstAttr::get(builder.getContext(), integer.getValue(), type) + ); + } + if (llvm::isa(resultType)) { + return builder.create(loc, fromAPInt(integer.getValue())); + } + if (resultType.isSignlessInteger(1)) { + return builder.create( + loc, integer.getValue().isZero() ? 0 : 1, resultType + ); + } + } else if (FeltConstAttr felt = llvm::dyn_cast(value)) { + return builder.create(loc, felt); + } + return nullptr; +} + +/// Rewrite `poly.const_read` uses in cloned bodies to concrete constants when their referenced +/// template parameter has been instantiated. class ClonedBodyConstReadOpPattern : public SymbolUserHelper< ClonedBodyConstReadOpPattern, ConstReadOp, IntegerAttr, FeltConstAttr> { @@ -331,6 +429,8 @@ class ClonedBodyConstReadOpPattern SymbolUserHelper; public: + /// Construct the const-read conversion pattern and collect delayed diagnostics in + /// `instantiationDiagnostics`. ClonedBodyConstReadOpPattern( TypeConverter &converter, MLIRContext *ctx, const DenseMap ¶mNameToInstantiatedValue, @@ -340,8 +440,11 @@ class ClonedBodyConstReadOpPattern : super(converter, ctx, /*patternBenefit=*/1, paramNameToInstantiatedValue), diagnostics(instantiationDiagnostics) {} + /// Use the referenced constant symbol as the lookup key. Attribute getNameAttr(ConstReadOp op) const override { return op.getConstNameAttr(); } + /// Replace an integer-backed template value with the constant op matching the converted result + /// type. LogicalResult handleRewrite( Attribute sym, ConstReadOp op, OpAdaptor, ConversionPatternRewriter &rewriter, IntegerAttr a ) const { @@ -352,25 +455,9 @@ class ClonedBodyConstReadOpPattern return op->emitOpError().append("could not convert result type ", origResTy); } - if (FeltType ty = llvm::dyn_cast(newResTy)) { - replaceOpWithNewOp( - rewriter, op, FeltConstAttr::get(getContext(), attrValue, ty) - ); - return success(); - } - - if (llvm::isa(newResTy)) { - replaceOpWithNewOp(rewriter, op, fromAPInt(attrValue)); - return success(); - } - if (newResTy.isSignlessInteger(1)) { // Treat 0 as false and any other value as true (but give a warning if it's not 1) - if (attrValue.isZero()) { - replaceOpWithNewOp(rewriter, op, false, newResTy); - return success(); - } - if (!attrValue.isOne()) { + if (!attrValue.isZero() && !attrValue.isOne()) { Location opLoc = op.getLoc(); Diagnostic diag(opLoc, DiagnosticSeverity::Warning); diag << "Interpreting non-zero value " << stringWithoutType(a) << " as true"; @@ -382,16 +469,21 @@ class ClonedBodyConstReadOpPattern << "\" for this call"; diagnostics.push_back(std::move(diag)); } - replaceOpWithNewOp(rewriter, op, true, newResTy); + } + if (Value replacement = materializeTemplateConstant(rewriter, op.getLoc(), newResTy, a)) { + rewriter.replaceOp(op, replacement); return success(); } return op->emitOpError().append("unexpected result type ", newResTy); } + /// Replace an already-felt template value with a felt constant. LogicalResult handleRewrite( Attribute, ConstReadOp op, OpAdaptor, ConversionPatternRewriter &rewriter, FeltConstAttr a ) const { - replaceOpWithNewOp(rewriter, op, a); + Value replacement = materializeTemplateConstant(rewriter, op.getLoc(), op.getType(), a); + assert(replacement && "felt template values must materialize as felt constants"); + rewriter.replaceOp(op, replacement); return success(); } }; @@ -401,8 +493,10 @@ class ClonedBodyConstReadOpPattern struct MatchFailureListener : public RewriterBase::Listener { bool hadFailure = false; - ~MatchFailureListener() override {} + /// Destroy the listener through the MLIR listener base class. + ~MatchFailureListener() override = default; + /// Convert match failures into reported diagnostics and remember that the pass must fail. void notifyMatchFailure(Location loc, function_ref reasonCallback) override { hadFailure = true; @@ -412,6 +506,8 @@ struct MatchFailureListener : public RewriterBase::Listener { } }; +/// Apply a greedy rewrite set, record whether it changed the module, and fail if any pattern +/// reported a hard match failure through `MatchFailureListener`. static LogicalResult applyAndFoldGreedily(ModuleOp modOp, ConversionTracker &tracker, RewritePatternSet &&patterns) { bool currStepModified = false; @@ -430,6 +526,63 @@ template bool isConcreteAttr(Attribute a) { return classifyAttrConcreteness(a, AllowStructParams) == AttrConcreteness::Concrete; } +/// Helper for applying template-parameter substitutions to attributes embedded in types. +class TemplateParamSubstitutions { + const DenseMap ¶mNameToValue; + +public: + /// Store the template-parameter binding map used by all substitution helpers. + explicit TemplateParamSubstitutions(const DenseMap &bindings) + : paramNameToValue(bindings) {} + + /// Return the bound value for `a` when present, otherwise return `a` unchanged. + Attribute lookupOrSelf(Attribute a) const { + auto res = paramNameToValue.find(a); + return (res != paramNameToValue.end()) ? res->second : a; + } + + /// Return true iff `nameAttr` has a concrete binding. + bool contains(Attribute nameAttr) const { return paramNameToValue.contains(nameAttr); } + + /// Replace a type variable with a concrete type binding when the binding is usable here. + Type convertTypeVarBinding(TypeVarType inputTy) const { + if (TypeAttr tyAttr = llvm::dyn_cast(lookupOrSelf(inputTy.getNameRef()))) { + Type convertedType = tyAttr.getValue(); + if (isConcreteType(convertedType)) { + return convertedType; + } + } + return inputTy; + } + + /// Substitute attributes, recursively converting nested `TypeAttr` payloads through `converter`. + SmallVector convertAttrs( + const TypeConverter &converter, ArrayRef attrs, bool *changed = nullptr + ) const { + SmallVector updated; + bool anyChanged = false; + for (Attribute attr : attrs) { + Attribute converted = attr; + if (TypeAttr tyAttr = dyn_cast(attr)) { + Type newTy = converter.convertType(tyAttr.getValue()); + if (newTy != tyAttr.getValue()) { + converted = TypeAttr::get(newTy); + } + } else { + converted = lookupOrSelf(attr); + } + anyChanged |= (converted != attr); + updated.push_back(converted); + } + if (changed != nullptr) { + *changed = anyChanged; + } + return updated; + } +}; + +/// Replace a callee rooted at a template parameter with the concrete struct callee named by that +/// parameter's instantiated type. static SymbolRefAttr convertCalleeSymRefs(SymbolRefAttr callee, const DenseMap ¶mNameToValue) { auto it = paramNameToValue.find(FlatSymbolRefAttr::get(callee.getRootReference())); @@ -452,6 +605,7 @@ convertCalleeSymRefs(SymbolRefAttr callee, const DenseMap return asSymbolRefAttr(newPieces); } +/// Rewrite all nested calls in `op` whose callee root names a concretized template parameter. static void convertCalleesInPlace(Operation *op, const DenseMap ¶mNameToValue) { op->walk([¶mNameToValue](CallOp callOp) { @@ -459,6 +613,7 @@ convertCalleesInPlace(Operation *op, const DenseMap ¶m }); } +/// Return true iff `op` calls a single-nested symbol rooted at a parameter of its parent template. static bool calleeReferencesTemplateParam(CallOp op) { SymbolRefAttr callee = op.getCalleeAttr(); if (!callee || callee.getNestedReferences().size() != 1) { @@ -555,6 +710,7 @@ evaluateTemplateExprs(TemplateOp templateOp, DenseMap &par ); } +/// Return true iff `op` no longer has a symbolic member table offset. static inline bool tableOffsetIsntSymbol(MemberReadOp op) { return !llvm::isa_and_present(op.getTableOffset().value_or(nullptr)); } @@ -566,6 +722,7 @@ class ClonedMemberReadOpPattern using super = SymbolUserHelper; public: + /// Construct the member-read conversion pattern for the active instantiation map. ClonedMemberReadOpPattern( TypeConverter &converter, MLIRContext *ctx, const DenseMap ¶mNameToInstantiatedValue @@ -573,10 +730,12 @@ class ClonedMemberReadOpPattern // benefit>0 so this applies instead of GeneralTypeReplacePattern : super(converter, ctx, /*patternBenefit=*/1, paramNameToInstantiatedValue) {} + /// Use the table-offset attribute as the lookup key. Attribute getNameAttr(MemberReadOp op) const override { return op.getTableOffset().value_or(nullptr); } + /// Replace a symbolic table offset with the concrete index value. LogicalResult handleRewrite( Attribute, MemberReadOp op, OpAdaptor, ConversionPatternRewriter &rewriter, IntegerAttr a ) const { @@ -587,6 +746,7 @@ class ClonedMemberReadOpPattern return success(); } + /// Emit a diagnostic for concrete template bindings that cannot index member tables. LogicalResult handleDefaultRewrite( Attribute, MemberReadOp op, OpAdaptor, ConversionPatternRewriter &, Attribute a ) const override { @@ -595,6 +755,7 @@ class ClonedMemberReadOpPattern ); } + /// Rewrite only member reads whose table offset is still a symbol. LogicalResult matchAndRewrite( MemberReadOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter ) const override { @@ -607,6 +768,21 @@ class ClonedMemberReadOpPattern } }; +/// Add the common constant and member-offset materialization patterns for a cloned body. +static void addClonedBodyMaterializationPatterns( + ConversionTarget &target, RewritePatternSet &patterns, TypeConverter &converter, + MLIRContext *ctx, const DenseMap ¶mNameToConcrete, + SmallVector &delayedDiagnostics +) { + target.addDynamicallyLegalOp([¶mNameToConcrete](ConstReadOp op) { + return !paramNameToConcrete.contains(op.getConstNameAttr()); + }); + patterns.add( + converter, ctx, paramNameToConcrete, delayedDiagnostics + ); + patterns.add(converter, ctx, paramNameToConcrete); +} + namespace Step1_InstantiateStructs { /// Implements cloning a `StructDefOp` for a specific instantiation site, using the concrete @@ -620,21 +796,18 @@ class StructCloner { class MappedTypeConverter : public TypeConverter { StructType origTy; StructType newTy; - const DenseMap ¶mNameToValue; - - inline Attribute convertIfPossible(Attribute a) const { - auto res = this->paramNameToValue.find(a); - return (res != this->paramNameToValue.end()) ? res->second : a; - } + TemplateParamSubstitutions substitutions; public: + /// Build a converter for a cloned struct body, replacing `originalType` with `newType` and + /// substituting any concretized template parameters. MappedTypeConverter( StructType originalType, StructType newType, /// Instantiated values for the parameter names in `originalType` const DenseMap ¶mNameToInstantiatedValue ) : TypeConverter(), origTy(originalType), newTy(newType), - paramNameToValue(paramNameToInstantiatedValue) { + substitutions(paramNameToInstantiatedValue) { addConversion([](Type inputTy) { return inputTy; }); @@ -647,14 +820,8 @@ class StructCloner { } // Check for replacement of parameter symbol names with concrete values if (ArrayAttr inputTyParams = inputTy.getParams()) { - SmallVector updated; - for (Attribute a : inputTyParams) { - if (TypeAttr ta = dyn_cast(a)) { - updated.push_back(TypeAttr::get(this->convertType(ta.getValue()))); - } else { - updated.push_back(convertIfPossible(a)); - } - } + SmallVector updated = + substitutions.convertAttrs(*this, inputTyParams.getValue()); return getStructTypeWithParams(inputTy.getNameRef(), inputTy.getContext(), updated); } // Otherwise, return the type unchanged @@ -665,10 +832,7 @@ class StructCloner { // Check for replacement of parameter symbol names with concrete values ArrayRef dimSizes = inputTy.getDimensionSizes(); if (!dimSizes.empty()) { - SmallVector updated; - for (Attribute a : dimSizes) { - updated.push_back(convertIfPossible(a)); - } + SmallVector updated = substitutions.convertAttrs(*this, dimSizes); return ArrayType::get(this->convertType(inputTy.getElementType()), updated); } // Otherwise, return the type unchanged @@ -676,21 +840,14 @@ class StructCloner { }); addConversion([this](TypeVarType inputTy) -> Type { - // Check for replacement of parameter symbol name with a concrete type - if (TypeAttr tyAttr = llvm::dyn_cast(convertIfPossible(inputTy.getNameRef()))) { - Type convertedType = tyAttr.getValue(); - // Use the new type unless it contains a TypeVarType because a TypeVarType from a - // different struct references a parameter name from that other struct, not from the - // current struct so the reference would be invalid. - if (isConcreteType(convertedType)) { - return convertedType; - } - } - return inputTy; + // Keep unresolved type variables from other templates because they reference names that + // are not valid in the current struct. + return substitutions.convertTypeVarBinding(inputTy); }); } }; + /// Clone `typeAtCaller` if at least one of its parameters is concrete at the current use site. FailureOr genClone(StructType typeAtCaller, ArrayRef typeAtCallerParams) { LLVM_DEBUG(llvm::dbgs() << "[StructCloner] attempting clone of " << typeAtCaller << '\n'); // Find the StructDefOp for the original StructType @@ -761,6 +918,7 @@ class StructCloner { // Clone the original struct. StructDefOp newStruct = origStruct.clone(); convertCalleesInPlace(newStruct, paramNameToConcrete); + Operation *insertedCloneRoot = nullptr; if (layout.remainingNames.empty()) { // FULL INSTANTIATION CASE // Set name of the new struct by prepending its name with instantiated template name. newStruct.setSymName( @@ -769,6 +927,7 @@ class StructCloner { // Insert 'newStruct' into the parent ModuleOp of the original TemplateOp. Use the // `SymbolTable::insert()` function so that the name will be made unique if necessary. symTables.getSymbolTable(parentModule).insert(newStruct, Block::iterator(parentTemplate)); + insertedCloneRoot = newStruct.getOperation(); // Drop the old template name from the list. typeAtCallerSymPieces.pop_back(); } else { // PARTIAL INSTANTIATION CASE @@ -793,6 +952,7 @@ class StructCloner { // `SymbolTable::insert()` function so that the name will be made unique if necessary. symTables.getSymbolTable(newTemplate).insert(newStruct); symTables.getSymbolTable(parentModule).insert(newTemplate, Block::iterator(parentTemplate)); + insertedCloneRoot = newTemplate.getOperation(); // Replace the old template name in the list with the new one (get template name after // symbol table insertion since it may be modified to make it unique). @@ -821,27 +981,26 @@ class StructCloner { MappedTypeConverter tyConv(typeAtDef, newStruct.getType(), paramNameToConcrete); ConversionTarget target = newConverterDefinedTarget(tyConv, ctx, tableOffsetIsntSymbol); - target.addDynamicallyLegalOp([¶mNameToConcrete](ConstReadOp op) { - // Legal if it's not in the map of concrete attribute instantiations - return !paramNameToConcrete.contains(op.getConstNameAttr()); - }); - RewritePatternSet patterns = newGeneralRewritePatternSet(tyConv, ctx, target); - patterns.add( - tyConv, ctx, paramNameToConcrete, tracker_.delayedDiagnosticSet(newLocalType) + addClonedBodyMaterializationPatterns( + target, patterns, tyConv, ctx, paramNameToConcrete, + tracker_.delayedDiagnosticSet(newRemoteType) ); - patterns.add(tyConv, ctx, paramNameToConcrete); if (failed(applyFullConversion(newStruct, target, std::move(patterns)))) { LLVM_DEBUG(llvm::dbgs() << "[StructCloner] instantiating body of struct failed \n"); + assert(insertedCloneRoot && "clone root must have been inserted before body conversion"); + insertedCloneRoot->erase(); return failure(); } return newRemoteType; } public: + /// Construct a cloner rooted at `root` and reporting modifications through `tracker`. StructCloner(ConversionTracker &tracker, ModuleOp root) : tracker_(tracker), rootMod(root), symTables() {} + /// Create a full or partial instantiated clone for `orig`, if `orig` has concrete parameters. FailureOr createInstantiatedClone(StructType orig) { LLVM_DEBUG(llvm::dbgs() << "[StructCloner] orig: " << orig << '\n'); if (ArrayAttr params = orig.getParams()) { @@ -851,8 +1010,10 @@ class StructCloner { return failure(); } + /// Re-enable diagnostics when a referenced struct definition cannot be found. void enableReportMissing() { reportMissing = true; } + /// Temporarily suppress missing-symbol diagnostics during speculative legality checks. void disableReportMissing() { reportMissing = false; } }; @@ -865,6 +1026,7 @@ class ParameterizedStructUseTypeConverter : public TypeConverter { friend DisableReportMissing; public: + /// Build a type converter that instantiates parameterized struct uses on demand. ParameterizedStructUseTypeConverter(ConversionTracker &tracker, ModuleOp root) : TypeConverter(), tracker_(tracker), cloner(tracker, root) { @@ -901,14 +1063,19 @@ class ParameterizedStructUseTypeConverter : public TypeConverter { } }; +/// Rewrite calls to struct `compute`/`constrain` functions after their struct types have been +/// instantiated. class CallStructFuncPattern : public OpConversionPattern { ConversionTracker &tracker_; public: + /// Construct the call rewrite pattern using the active type converter and tracker. CallStructFuncPattern(TypeConverter &converter, MLIRContext *ctx, ConversionTracker &tracker) // benefit>0 so this applies instead of CallOpClassReplacePattern : OpConversionPattern(converter, ctx, /*benefit=*/1), tracker_(tracker) {} + /// Replace a call with converted result types and, when needed, a callee rooted at the + /// instantiated struct type. LogicalResult matchAndRewrite( CallOp op, OpAdaptor adapter, ConversionPatternRewriter &rewriter ) const override { @@ -952,13 +1119,15 @@ class CallStructFuncPattern : public OpConversionPattern { } }; -// This one ensures MemberDefOp types are converted even if there are no reads/writes to them. +/// Ensure `struct.member` types are converted even if no read/write pattern visits them. class MemberDefOpPattern : public OpConversionPattern { public: + /// Construct the member definition conversion pattern. MemberDefOpPattern(TypeConverter &converter, MLIRContext *ctx, ConversionTracker &) // benefit>0 so this applies instead of GeneralTypeReplacePattern : OpConversionPattern(converter, ctx, /*benefit=*/1) {} + /// Update the member definition type when the active type converter changes it. LogicalResult matchAndRewrite( MemberDefOp op, OpAdaptor /*adapter*/, ConversionPatternRewriter &rewriter ) const override { @@ -974,26 +1143,58 @@ class MemberDefOpPattern : public OpConversionPattern { } }; +/// Convert nondeterministic result types when a shared witness is refined to an instantiated type. +class NonDetOpPattern : public OpConversionPattern { +public: + /// Construct the nondeterministic value conversion pattern. + NonDetOpPattern(TypeConverter &converter, MLIRContext *ctx) + : OpConversionPattern(converter, ctx, /*benefit=*/1) {} + + /// Rebuild the op with the converted result type. + LogicalResult + matchAndRewrite(NonDetOp op, OpAdaptor, ConversionPatternRewriter &rewriter) const override { + Type newType = getTypeConverter()->convertType(op.getType()); + if (!newType) { + return op->emitError("Could not convert Op result type."); + } + if (newType == op.getType()) { + return failure(); + } + NonDetOp newOp = rewriter.create(op.getLoc(), newType); + newOp->setDiscardableAttrs(op->getDiscardableAttrDictionary()); + rewriter.replaceOp(op, newOp.getResult()); + return success(); + } +}; + /// Disables reporting of missing struct symbols during legality checks to avoid showing error /// diagnostics that are not actually errors. class DisableReportMissing : public LegalityCheckCallback { ParameterizedStructUseTypeConverter &tyConv; public: + /// Tie the callback to the converter whose cloner should suppress lookup diagnostics. explicit DisableReportMissing(ParameterizedStructUseTypeConverter &tc) : tyConv(tc) {} + /// Suppress missing-symbol diagnostics before a speculative legality check begins. void checkStarted() override { tyConv.cloner.disableReportMissing(); } + /// Re-enable missing-symbol diagnostics after the speculative legality check finishes. void checkEnded(bool) override { tyConv.cloner.enableReportMissing(); } }; +/// Run struct instantiation and call/member rewrites for the current module. LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { MLIRContext *ctx = modOp.getContext(); ParameterizedStructUseTypeConverter tyConv(tracker, modOp); DisableReportMissing drm(tyConv); ConversionTarget target = newConverterDefinedTargetWithCallback<>(tyConv, ctx, drm); + target.addDynamicallyLegalOp([&tyConv](Operation *op) { + return defaultLegalityCheck(tyConv, op); + }); RewritePatternSet patterns = newGeneralRewritePatternSet(tyConv, ctx, target); patterns.add(tyConv, ctx, tracker); + patterns.add(tyConv, ctx); return applyPartialConversion(modOp, target, std::move(patterns)); } @@ -1030,37 +1231,23 @@ namespace Step2_InstantiateFunctions { /// ArrayType/StructType parameters with their concrete values determined by unification. class FuncInstTypeConverter : public TypeConverter { DenseMap paramNameToValue; - - Attribute convertIfPossible(Attribute a) const { - auto res = paramNameToValue.find(a); - return (res != paramNameToValue.end()) ? res->second : a; - } + TemplateParamSubstitutions substitutions; public: + /// Build the function-instantiation type converter from concrete template bindings. explicit FuncInstTypeConverter(DenseMap paramNameToConcrete) - : TypeConverter(), paramNameToValue(std::move(paramNameToConcrete)) { + : TypeConverter(), paramNameToValue(std::move(paramNameToConcrete)), + substitutions(paramNameToValue) { addConversion([](Type t) { return t; }); addConversion([this](TypeVarType inputTy) -> Type { - if (TypeAttr tyAttr = llvm::dyn_cast(convertIfPossible(inputTy.getNameRef()))) { - Type convertedType = tyAttr.getValue(); - if (isConcreteType(convertedType)) { - return convertedType; - } - } - return inputTy; + return substitutions.convertTypeVarBinding(inputTy); }); addConversion([this](ArrayType inputTy) { - SmallVector updated; bool changed = false; - for (Attribute a : inputTy.getDimensionSizes()) { - Attribute converted = convertIfPossible(a); - updated.push_back(converted); - if (converted != a) { - changed = true; - } - } + SmallVector updated = + substitutions.convertAttrs(*this, inputTy.getDimensionSizes(), &changed); Type newElemTy = this->convertType(inputTy.getElementType()); if (!changed && newElemTy == inputTy.getElementType()) { return inputTy; @@ -1072,26 +1259,9 @@ class FuncInstTypeConverter : public TypeConverter { addConversion([this](StructType inputTy) -> StructType { if (ArrayAttr params = inputTy.getParams()) { - SmallVector updated; bool changed = false; - for (Attribute a : params) { - if (TypeAttr ta = dyn_cast(a)) { - Type newTy = this->convertType(ta.getValue()); - if (newTy != ta.getValue()) { - updated.push_back(TypeAttr::get(newTy)); - changed = true; - continue; - } - } else { - Attribute converted = convertIfPossible(a); - if (converted != a) { - updated.push_back(converted); - changed = true; - continue; - } - } - updated.push_back(a); - } + SmallVector updated = + substitutions.convertAttrs(*this, params.getValue(), &changed); if (changed) { return getStructTypeWithParams(inputTy.getNameRef(), inputTy.getContext(), updated); } @@ -1100,6 +1270,7 @@ class FuncInstTypeConverter : public TypeConverter { }); } + /// Convert an attribute that may contain a type or a direct template-parameter reference. Attribute convertAttr(Attribute attr) const { if (TypeAttr tyAttr = llvm::dyn_cast(attr)) { Type convertedTy = convertType(tyAttr.getValue()); @@ -1107,15 +1278,18 @@ class FuncInstTypeConverter : public TypeConverter { return TypeAttr::get(convertedTy); } } - return convertIfPossible(attr); + return substitutions.lookupOrSelf(attr); } - bool containsParam(Attribute nameAttr) const { return paramNameToValue.contains(nameAttr); } + /// Return true iff the given template parameter has a concrete binding in this converter. + bool containsParam(Attribute nameAttr) const { return substitutions.contains(nameAttr); } + + /// Return the underlying template-parameter binding map. const DenseMap &getParamMap() const { return paramNameToValue; } }; /// Return the callee-side unification-derived value for a template parameter, if any. -inline static std::optional +static inline std::optional inferUnifiedParam(const UnificationMap &unifyResult, SymbolRefAttr paramName) { auto it = unifyResult.find({paramName, Side::RHS}); return (it == unifyResult.end()) ? std::nullopt : std::make_optional(it->second); @@ -1123,7 +1297,7 @@ inferUnifiedParam(const UnificationMap &unifyResult, SymbolRefAttr paramName) { /// Emit the match failure used when an inferred instantiation violates a template parameter's /// declared type restriction. -inline static LogicalResult failIncompatibleInferredParam( +static inline LogicalResult failIncompatibleInferredParam( CallOp op, PatternRewriter &rewriter, FlatSymbolRefAttr paramName, TemplateParamOp paramOp ) { LLVM_DEBUG( @@ -1147,12 +1321,23 @@ class WildcardTypeBodyInferer final { SmallVector> activeInferences_; public: + /// Construct a body inferer over the current symbol tables and known concrete bindings. WildcardTypeBodyInferer( SymbolTableCollection &symTables, const DenseMap ¶mNameToConcrete ) : symTables_(symTables), paramNameToConcrete_(paramNameToConcrete) {} + /// Search `func` for a concrete value that can resolve `paramName`. std::optional infer(FuncDefOp func, FlatSymbolRefAttr paramName) { + return infer(func, paramName, paramNameToConcrete_); + } + +private: + /// Search `func` using concrete bindings expressed in that function's template scope. + std::optional infer( + FuncDefOp func, FlatSymbolRefAttr paramName, + const DenseMap ¶mNameToConcrete + ) { if (llvm::any_of(activeInferences_, [&](const auto &e) { return e.first == func.getOperation() && e.second == paramName; })) { @@ -1160,7 +1345,7 @@ class WildcardTypeBodyInferer final { } activeInferences_.emplace_back(func.getOperation(), paramName); - FuncInstTypeConverter tyConv((paramNameToConcrete_)); + FuncInstTypeConverter tyConv(paramNameToConcrete); std::optional inferred; bool ambiguous = false; @@ -1230,7 +1415,9 @@ class WildcardTypeBodyInferer final { } continue; } - if (std::optional candidate = infer(nestedTgt, nestedTvar.getNameRef())) { + DenseMap nestedParamNameToConcrete = + getNestedConcreteBindings(nestedCall, nestedTgt, nestedTemplate, tyConv); + if (auto candidate = infer(nestedTgt, nestedTvar.getNameRef(), nestedParamNameToConcrete)) { WalkResult candidateResult = noteCandidate(*candidate); if (candidateResult.wasInterrupted()) { return candidateResult; @@ -1247,7 +1434,46 @@ class WildcardTypeBodyInferer final { return inferred; } -private: + /// Map concrete bindings for a nested callee into the nested template's parameter scope. + /// + /// The current converter's keys name parameters of the enclosing template. Reusing it while + /// walking a nested callee would therefore incorrectly bind same-named nested parameters. + static DenseMap getNestedConcreteBindings( + CallOp nestedCall, FuncDefOp nestedTgt, TemplateOp nestedTemplate, + const FuncInstTypeConverter &enclosingTyConv + ) { + DenseMap bindings; + auto nestedParams = nestedTemplate.getConstOps(); + ArrayAttr callParams = nestedCall.getTemplateParamsAttr(); + + if (!isNullOrEmpty(callParams)) { + for (auto [paramOp, arg] : llvm::zip_equal(nestedParams, callParams.getValue())) { + Attribute value = enclosingTyConv.convertAttr(arg); + if (isConcreteAttr(value)) { + bindings[FlatSymbolRefAttr::get(paramOp.getSymNameAttr())] = value; + } + } + return bindings; + } + + FailureOr unifyResult = + nestedCall.unifyTypeSignature(nestedTgt.getFunctionType()); + if (failed(unifyResult)) { + return bindings; + } + for (TemplateParamOp paramOp : nestedParams) { + auto paramName = FlatSymbolRefAttr::get(paramOp.getSymNameAttr()); + if (std::optional value = inferUnifiedParam(*unifyResult, paramName)) { + Attribute convertedValue = enclosingTyConv.convertAttr(*value); + if (isConcreteAttr(convertedValue)) { + bindings[paramName] = convertedValue; + } + } + } + return bindings; + } + + /// Infer a nested callee parameter value from the nested call's explicit template arguments. std::optional inferFromExplicitNestedCallParams( CallOp nestedCall, TemplateOp nestedTemplate, FlatSymbolRefAttr nestedParamName, const FuncInstTypeConverter &tyConv @@ -1276,6 +1502,7 @@ class ClonedBodyArrayReadOpPattern final : public OpConversionPattern::OpConversionPattern; + /// Replace a scalar element read with an array extract when conversion makes the result an array. LogicalResult matchAndRewrite( ReadArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter ) const override { @@ -1283,9 +1510,10 @@ class ClonedBodyArrayReadOpPattern final : public OpConversionPattern(newResultTy)) { return failure(); } - replaceOpWithNewOp( - rewriter, op, newResultTy, adaptor.getArrRef(), adaptor.getIndices() + ExtractArrayOp extractOp = rewriter.replaceOpWithNewOp( + op, newResultTy, adaptor.getArrRef(), adaptor.getIndices() ); + extractOp->setDiscardableAttrs(op->getDiscardableAttrDictionary()); return success(); } }; @@ -1296,6 +1524,7 @@ class ClonedBodyArrayWriteOpPattern final : public OpConversionPattern::OpConversionPattern; + /// Replace a scalar element write with an array insert when conversion makes the value an array. LogicalResult matchAndRewrite( WriteArrayOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter ) const override { @@ -1310,25 +1539,20 @@ class ClonedBodyArrayWriteOpPattern final : public OpConversionPattern ¶mNameToConcrete ) { MLIRContext *ctx = op.getContext(); FuncInstTypeConverter tyConv(paramNameToConcrete); - ConversionTarget target = newConverterDefinedTarget<>(tyConv, ctx, tableOffsetIsntSymbol); - target.addDynamicallyLegalOp([&tyConv](ConstReadOp p) { - // Legal if it's not in the map of concrete attribute instantiations - return !tyConv.containsParam(p.getConstNameAttr()); - }); SmallVector delayedDiagnostics; + ConversionTarget target = newConverterDefinedTarget<>(tyConv, ctx, tableOffsetIsntSymbol); RewritePatternSet bodyPatterns = newGeneralRewritePatternSet(tyConv, ctx, target); - bodyPatterns.add( - tyConv, ctx, tyConv.getParamMap(), delayedDiagnostics + addClonedBodyMaterializationPatterns( + target, bodyPatterns, tyConv, ctx, tyConv.getParamMap(), delayedDiagnostics ); bodyPatterns.add(tyConv, ctx); - bodyPatterns.add(tyConv, ctx, paramNameToConcrete); if (failed(applyFullConversion(newFunc, target, std::move(bodyPatterns)))) { return failure(); } @@ -1342,13 +1566,123 @@ static LogicalResult applyBodyConversions( return failure(res.wasInterrupted()); } +/// Copy unresolved template expressions referenced by `newFuncs` into their partially-instantiated +/// parent template. Reads of concrete parameters within the copied expressions are materialized so +/// the new template contains no references to parameters that it does not preserve. +static LogicalResult copyReferencedTemplateExprs( + TemplateOp parentTemplate, Block &newTemplateBody, ArrayRef newFuncs, + const DenseMap ¶mNameToConcrete +) { + DenseSet referencedExprNames; + auto collectExprRef = [&referencedExprNames, ¶mNameToConcrete](FlatSymbolRefAttr name) { + if (!paramNameToConcrete.contains(name)) { + referencedExprNames.insert(name.getAttr()); + } + }; + + // Expression references can occur in a function signature, in type-bearing operands/results, + // or in arbitrary attributes as well as in poly.read_const operations. In particular, a + // signature-only reference must keep its defining expression in a partial template even if the + // function body does not read that expression. + for (FuncDefOp newFunc : newFuncs) { + newFunc.walk([&](Operation *nestedOp) { + auto collectTypeRefs = [&collectExprRef](Type type) { type.walk(collectExprRef); }; + for (Type type : nestedOp->getOperandTypes()) { + collectTypeRefs(type); + } + for (Type type : nestedOp->getResultTypes()) { + collectTypeRefs(type); + } + nestedOp->getAttrDictionary().walk(collectExprRef); + }); + newFunc.walk([&collectExprRef](ConstReadOp readOp) { + collectExprRef(readOp.getConstNameAttr()); + }); + } + + for (TemplateExprOp expr : parentTemplate.getConstOps()) { + if (!referencedExprNames.contains(expr.getSymNameAttr())) { + continue; + } + + auto clonedExpr = llvm::cast(expr->clone()); + SmallVector concreteReads; + clonedExpr.walk([&](ConstReadOp readOp) { + if (paramNameToConcrete.contains(readOp.getConstNameAttr())) { + concreteReads.push_back(readOp); + } + }); + for (ConstReadOp readOp : concreteReads) { + Attribute value = paramNameToConcrete.lookup(readOp.getConstNameAttr()); + OpBuilder builder(readOp); + Value replacement = + materializeTemplateConstant(builder, readOp.getLoc(), readOp.getType(), value); + if (!replacement) { + clonedExpr->erase(); + return failure(); + } + readOp.replaceAllUsesWith(replacement); + readOp.erase(); + } + newTemplateBody.push_back(clonedExpr); + } + return success(); +} + +/// Clone every free function in `parentTemplate` reached through a call in the cloned functions. +/// Rewriting these calls to the new template keeps partial instantiations self-contained. +static SmallVector copyReferencedTemplateSiblingFuncs( + TemplateOp parentTemplate, TemplateOp newTemplate, FuncDefOp originalFunc, FuncDefOp newFunc, + SymbolTableCollection &symTables, const DenseMap ¶mNameToConcrete +) { + DenseMap cloned; + SmallVector copiedFuncs {newFunc}; + cloned[originalFunc] = newFunc; + SymbolTable &parentSymbols = symTables.getSymbolTable(parentTemplate); + + for (size_t i = 0; i < copiedFuncs.size(); ++i) { + FuncDefOp current = copiedFuncs[i]; + current.walk([&](CallOp nestedCall) { + SymbolRefAttr callee = nestedCall.getCalleeAttr(); + if (callee.getRootReference() != parentTemplate.getSymName() || + callee.getNestedReferences().size() != 1) { + return; + } + auto sibling = dyn_cast_or_null( + parentSymbols.lookup(callee.getNestedReferences().front().getAttr()) + ); + if (!sibling || sibling->getParentOp() != parentTemplate) { + return; + } + + FuncDefOp clonedSibling; + if (auto found = cloned.find(sibling); found != cloned.end()) { + clonedSibling = found->second; + } else { + clonedSibling = sibling.clone(); + convertCalleesInPlace(clonedSibling, paramNameToConcrete); + cloned[sibling] = clonedSibling; + copiedFuncs.push_back(clonedSibling); + } + nestedCall.setCalleeAttr( + SymbolRefAttr::get( + newTemplate.getSymNameAttr(), {FlatSymbolRefAttr::get(clonedSibling.getSymNameAttr())} + ) + ); + }); + } + return copiedFuncs; +} + class InstantiateFuncAtCallOp final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the function-instantiation pattern. InstantiateFuncAtCallOp(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx), tracker_(tracker) {} + /// Instantiate the target function or template at a call site and rewrite the callee reference. LogicalResult matchAndRewrite(CallOp op, PatternRewriter &rewriter) const override { LLVM_DEBUG(llvm::dbgs() << "[InstantiateFuncAtCallOp] op: " << op << '\n'); @@ -1394,10 +1728,10 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { // Maps template parameter symbols to the instantiation value at the call site. DenseMap paramNameToConcrete; - if (failed(collectConcreteTemplateParams( - op, rewriter, symTables, callTgt, parentTemplate, unifyResult.value(), - paramNameToConcrete - ))) { + auto collectRes = collectConcreteTemplateParams( + op, rewriter, symTables, callTgt, parentTemplate, unifyResult.value(), paramNameToConcrete + ); + if (failed(collectRes)) { return failure(); } @@ -1417,22 +1751,25 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { ModuleOp parentModule = getParentOfType(parentTemplate); assert(parentModule && "TemplateOp must be nested in a ModuleOp"); - SymbolRefAttr originalCalleeAttr = op.getCalleeAttr(); FailureOr newCalleeAttr = - layout.remainingNames.empty() - ? instantiateFully( - op, rewriter, symTables, callTgt, parentTemplate, parentModule, - layout.templateNameWithAttrs, paramNameToConcrete - ) - : instantiatePartially( - op, rewriter, symTables, callTgt, parentTemplate, parentModule, layout, - paramNameToConcrete, tracker_ - ); + layout.remainingNames.empty() ? instantiateFully( + op, rewriter, symTables, callTgt, parentTemplate, + parentModule, layout.templateNameWithAttrs, + layout.concreteParamKey, paramNameToConcrete, tracker_ + ) + : instantiatePartially( + op, rewriter, symTables, callTgt, parentTemplate, + parentModule, layout, paramNameToConcrete, tracker_ + ); if (failed(newCalleeAttr)) { return failure(); } - tracker_.recordInstantiation(originalCalleeAttr); + FailureOr originalCalleePath = getPathFromTopRoot(callTgt); + if (failed(originalCalleePath)) { + return failure(); + } + tracker_.recordInstantiation(*originalCalleePath); // Update the CallOp to point to the instantiated function and mark the module as modified. rewriter.modifyOpInPlace(op, [&op, &newCalleeAttr, &layout]() { @@ -1520,12 +1857,12 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { // instantiation is valid, except for the size check because that cannot change. assert((callParams.size() == llvm::range_size(realParams)) && "per CallOpVerifier"); if (failed(op.verifyTemplateParamCompatibility(realParams))) { - return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) { + return rewriter.notifyMatchFailure(op, [](Diagnostic &diag) { diag.append("incompatible with specified param type(s)"); }); } if (failed(op.verifyTemplateParamsMatchInferred(realParams, unifyResult))) { - return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) { + return rewriter.notifyMatchFailure(op, [](Diagnostic &diag) { diag.append("incompatible with inferred param value(s)"); }); } @@ -1578,54 +1915,67 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { return success(); } + /// Return the full-instantiation callee spelling that is valid at this call site. + static SymbolRefAttr + buildFullInstantiationCalleeForCall(CallOp op, FlatSymbolRefAttr newFuncName) { + // Callee: drop template & original function names, add the new module-level function name. + // Original: @[prefix...]::@TemplateName::@funcName + // New: @[prefix...]::@newFuncName + SmallVector symPieces = getPieces(op.getCalleeAttr()); + assert(symPieces.size() >= 2 && "callee must include at least template and function names"); + symPieces.pop_back(); // remove original function name + symPieces.pop_back(); // remove template name + symPieces.push_back(newFuncName); + return asSymbolRefAttr(symPieces); + } + /// Create or reuse a fully-instantiated clone in the parent module and return the rewritten /// module-level callee reference. static FailureOr instantiateFully( CallOp op, PatternRewriter &rewriter, SymbolTableCollection &symTables, FuncDefOp callTgt, TemplateOp parentTemplate, ModuleOp parentModule, StringRef templateNameWithAttrs, - const DenseMap ¶mNameToConcrete + ArrayAttr concreteParamKey, const DenseMap ¶mNameToConcrete, + ConversionTracker &tracker ) { - MLIRContext *ctx = op.getContext(); - std::string newFuncName = - (mlir::Twine(templateNameWithAttrs) + "_" + callTgt.getSymName()).str(); - StringRef actualNewFuncName = newFuncName; - if (!symTables.getSymbolTable(parentModule).lookup(newFuncName)) { - FuncDefOp newFunc = callTgt.clone(); - newFunc.setSymName(newFuncName); - convertCalleesInPlace(newFunc, paramNameToConcrete); - // Insert before the TemplateOp; symbol table may adjust the name to ensure uniqueness. - symTables.getSymbolTable(parentModule).insert(newFunc, Block::iterator(parentTemplate)); - actualNewFuncName = newFunc.getSymName(); + if (auto cachedPath = tracker.getFullFuncInstantiation(callTgt, concreteParamKey)) { LLVM_DEBUG( - llvm::dbgs() << "[InstantiateFuncAtCallOp] created full instantiation function: " - << actualNewFuncName << '\n' + llvm::dbgs() << "[InstantiateFuncAtCallOp] reusing full instantiation function: " + << *cachedPath << '\n' ); - if (failed(applyBodyConversions(op, newFunc, paramNameToConcrete))) { - LLVM_DEBUG( - llvm::dbgs() << "[InstantiateFuncAtCallOp] body conversion failed for " - << actualNewFuncName << '\n' - ); - newFunc->erase(); - return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) { - diag.append("failure while creating instantiated function '", actualNewFuncName, '\''); - }); - } - } else { + return buildFullInstantiationCalleeForCall(op, getPieces(*cachedPath).back()); + } + + std::string newFuncName; + llvm::raw_string_ostream(newFuncName) << templateNameWithAttrs << '_' << callTgt.getSymName(); + FuncDefOp newFunc = callTgt.clone(); + newFunc.setSymName(newFuncName); + convertCalleesInPlace(newFunc, paramNameToConcrete); + // Insert before the TemplateOp; symbol table may adjust the name to avoid existing symbols. + symTables.getSymbolTable(parentModule).insert(newFunc, Block::iterator(parentTemplate)); + StringAttr actualNewFuncName = newFunc.getSymNameAttr(); + LLVM_DEBUG( + llvm::dbgs() << "[InstantiateFuncAtCallOp] created full instantiation function: " + << actualNewFuncName << '\n' + ); + if (failed(applyBodyConversions(op, newFunc, paramNameToConcrete))) { LLVM_DEBUG( - llvm::dbgs() << "[InstantiateFuncAtCallOp] reusing full instantiation function: " + llvm::dbgs() << "[InstantiateFuncAtCallOp] body conversion failed for " << actualNewFuncName << '\n' ); + newFunc->erase(); + return rewriter.notifyMatchFailure(op, [&actualNewFuncName](Diagnostic &diag) { + diag.append("failure while creating instantiated function ", actualNewFuncName); + }); } - // Callee: drop template & original function names, add the new module-level function name. - // Original: @[prefix...]::@TemplateName::@funcName - // New: @[prefix...]::@newFuncName - SmallVector symPieces = getPieces(op.getCalleeAttr()); - assert(symPieces.size() >= 2 && "callee must include at least template and function names"); - symPieces.pop_back(); // remove original function name - symPieces.pop_back(); // remove template name - symPieces.push_back(FlatSymbolRefAttr::get(StringAttr::get(ctx, actualNewFuncName))); - return asSymbolRefAttr(symPieces); + FailureOr newCalleeAttr = getPathFromTopRoot(newFunc); + if (failed(newCalleeAttr)) { + return failure(); + } + tracker.recordFullFuncInstantiation(callTgt, concreteParamKey, *newCalleeAttr); + return buildFullInstantiationCalleeForCall( + op, FlatSymbolRefAttr::get(newFunc.getSymNameAttr()) + ); } /// Create or reuse a partially-instantiated template that preserves the remaining non-concrete @@ -1670,20 +2020,51 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { // Clone and partially convert the function (concretize only the concrete params). FuncDefOp newFunc = callTgt.clone(); convertCalleesInPlace(newFunc, paramNameToConcrete); + SmallVector copiedFuncs = copyReferencedTemplateSiblingFuncs( + parentTemplate, newTemplate, callTgt, newFunc, symTables, paramNameToConcrete + ); + auto copyRes = copyReferencedTemplateExprs( + parentTemplate, newTemplateBody, copiedFuncs, paramNameToConcrete + ); + if (failed(copyRes)) { + for (FuncDefOp copiedFunc : copiedFuncs) { + copiedFunc->erase(); + } + newTemplate->erase(); + return rewriter.notifyMatchFailure(op, "failure while copying template expressions"); + } // Insert before body conversion so nested concrete callees verify from the root module. Use // SymbolTable::insert() so both physical symbol names are unique if necessary. - symTables.getSymbolTable(newTemplate).insert(newFunc); + for (FuncDefOp copiedFunc : copiedFuncs) { + symTables.getSymbolTable(newTemplate).insert(copiedFunc); + } + StringAttr provisionalTemplateName = newTemplate.getSymNameAttr(); symTables.getSymbolTable(parentModule).insert(newTemplate, Block::iterator(parentTemplate)); - if (failed(applyBodyConversions(op, newFunc, paramNameToConcrete))) { - std::string newFuncName = newFunc.getSymName().str(); + if (newTemplate.getSymNameAttr() != provisionalTemplateName) { + for (FuncDefOp copiedFunc : copiedFuncs) { + copiedFunc.walk([&](CallOp nestedCall) { + SymbolRefAttr callee = nestedCall.getCalleeAttr(); + if (callee.getRootReference() == provisionalTemplateName) { + nestedCall.setCalleeAttr( + SymbolRefAttr::get(newTemplate.getSymNameAttr(), callee.getNestedReferences()) + ); + } + }); + } + } + for (FuncDefOp copiedFunc : copiedFuncs) { + if (succeeded(applyBodyConversions(op, copiedFunc, paramNameToConcrete))) { + continue; + } + StringAttr newFuncName = copiedFunc.getSymNameAttr(); LLVM_DEBUG( llvm::dbgs() << "[InstantiateFuncAtCallOp] body conversion failed for " << newFuncName << '\n' ); newTemplate->erase(); - return rewriter.notifyMatchFailure(op, [&](Diagnostic &diag) { - diag.append("failure while creating instantiated function '", newFuncName, '\''); + return rewriter.notifyMatchFailure(op, [&newFuncName](Diagnostic &diag) { + diag.append("failure while creating instantiated function ", newFuncName); }); } @@ -1706,6 +2087,7 @@ class InstantiateFuncAtCallOp final : public OpRewritePattern { } }; +/// Run function instantiation patterns once over the module. LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { MLIRContext *ctx = modOp.getContext(); RewritePatternSet patterns(ctx); @@ -1725,6 +2107,7 @@ class LoopUnrollPattern : public OpRewritePattern { public: using OpRewritePattern::OpRewritePattern; + /// Fully unroll loop-like ops whose trip count is statically known. LogicalResult matchAndRewrite(OpClass loopOp, PatternRewriter &rewriter) const override { if (auto maybeConstant = getConstantTripCount(loopOp)) { uint64_t tripCount = *maybeConstant; @@ -1753,6 +2136,7 @@ class LoopUnrollPattern : public OpRewritePattern { } }; +/// Run loop unrolling for supported loop dialects. LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { MLIRContext *ctx = modOp.getContext(); RewritePatternSet patterns(ctx); @@ -1765,8 +2149,10 @@ LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { namespace Step4_InstantiateAffineMaps { -// Adapted from `mlir::getConstantIntValues()` but that one failed in CI for an unknown reason. This -// version uses a basic loop instead of llvm::map_to_vector(). +/// Return constant integer values for all fold results, if every fold result is constant. +/// +/// Adapted from `mlir::getConstantIntValues()` but that one failed in CI for an unknown reason. +/// This version uses a basic loop instead of llvm::map_to_vector(). std::optional> getConstantIntValues(ArrayRef ofrs) { SmallVector res; for (OpFoldResult ofr : ofrs) { @@ -1779,25 +2165,36 @@ std::optional> getConstantIntValues(ArrayRef return res; } +/// Folds affine-map parameters using the map operands supplied at an instantiation site. struct AffineMapFolder { + /// Inputs that describe affine-map operands and the parameter list being folded. struct Input { + /// Operand groups corresponding to affine-map parameters. OperandRangeRange mapOpGroups; + /// Number of dimensions in each operand group. DenseI32ArrayAttr dimsPerGroup; + /// Parameter list containing affine maps and non-map attributes. ArrayRef paramsOfStructTy; }; + /// Outputs after replacing foldable affine-map parameters with concrete attributes. struct Output { + /// Operand groups for affine maps that could not be folded. SmallVector> mapOpGroups; + /// Dimension counts corresponding to remaining map operand groups. SmallVector dimsPerGroup; + /// Parameter list with folded values substituted where possible. SmallVector paramsOfStructTy; }; - static inline SmallVector getConvertedMapOpGroups(Output out) { + /// Convert owned output operand groups into `ValueRange` views for op builders. + static inline SmallVector getConvertedMapOpGroups(const Output &out) { return llvm::map_to_vector(out.mapOpGroups, [](const SmallVector &grp) { return ValueRange(grp); }); } + /// Fold any affine-map attributes in `in.paramsOfStructTy` whose operands are all constants. static LogicalResult fold(PatternRewriter &rewriter, const Input &in, Output &out, Operation *op, const char *aspect) { if (in.mapOpGroups.empty()) { @@ -1865,7 +2262,7 @@ struct AffineMapFolder { out.mapOpGroups.emplace_back(currMapOps); out.dimsPerGroup.push_back(in.dimsPerGroup[idx - 1]); // idx was already incremented } - // If not affine and foldable, preserve the original + // If not affine, preserve the original. out.paramsOfStructTy.push_back(sizeAttr); } assert(idx == in.mapOpGroups.size() && "all affine_map not processed"); @@ -1884,9 +2281,11 @@ class InstantiateAtCreateArrayOp final : public OpRewritePattern ConversionTracker &tracker_; public: + /// Construct the array-creation affine-map instantiation pattern. InstantiateAtCreateArrayOp(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx), tracker_(tracker) {} + /// Rewrite `array.new` when affine-map dimensions can be folded to concrete sizes. LogicalResult matchAndRewrite(CreateArrayOp op, PatternRewriter &rewriter) const override { ArrayType oldResultType = op.getType(); @@ -1922,9 +2321,11 @@ class InstantiateAtCallOpCompute final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the struct-compute call result instantiation pattern. InstantiateAtCallOpCompute(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx), tracker_(tracker) {} + /// Refine the result type of calls to struct `compute` functions when parameters become known. LogicalResult matchAndRewrite(CallOp op, PatternRewriter &rewriter) const override { if (!op.calleeIsStructCompute()) { // this pattern only applies when the callee is "compute()" within a struct @@ -2085,6 +2486,7 @@ class InstantiateAtCallOpCompute final : public OpRewritePattern { } }; +/// Run affine-map and target-type instantiation over arrays and struct `compute` calls. LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { MLIRContext *ctx = modOp.getContext(); RewritePatternSet patterns(ctx); @@ -2098,37 +2500,2712 @@ LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { } // namespace Step4_InstantiateAffineMaps -namespace Step5_PropagateTypes { +/// Rebuild a struct constraint call with its witness specialized to a concrete struct type. +static CallOp replaceStructConstraintCallWithSpecializedWitness( + PatternRewriter &rewriter, CallOp call, Value specializedWitness +) { + SmallVector args(call.getArgOperands()); + args.front() = specializedWitness; + StructType structType = llvm::cast(specializedWitness.getType()); + SymbolRefAttr callee = + appendLeaf(structType.getNameRef(), call.getCalleeAttr().getLeafReference()); + CallOp newCall = replaceOpWithNewOp( + rewriter, call, call.getResultTypes(), callee, + CallOp::toVectorOfValueRange(call.getMapOperands()), call.getNumDimsPerMapAttr(), args + ); + newCall->setDiscardableAttrs(call->getDiscardableAttrDictionary()); + return newCall; +} -/// Update the array element type by looking at the values stored into it from uses. -class UpdateNewArrayElemFromWrite final : public OpRewritePattern { - ConversionTracker &tracker_; +namespace Step5_ScalarizeHeterogeneousArrays { -public: - UpdateNewArrayElemFromWrite(MLIRContext *ctx, ConversionTracker &tracker) - : OpRewritePattern(ctx, 3), tracker_(tracker) {} +/// Information about a local array allocation that can be replaced with the final values of its +/// statically-known elements. +/// +/// This pass only scalarizes arrays after loop unrolling and affine-map instantiation have exposed +/// all element indices and value types. The candidate array must have either an initializer value +/// or exactly one write for every static element index, and all direct reads/member writes must +/// happen after every explicit element write. Those restrictions avoid imposing a new memory +/// semantics for partially initialized arrays, repeated writes, dynamic indices, or +/// branch-sensitive updates. +struct ScalarizedArrayInfo { + /// The array allocation being removed. + CreateArrayOp createOp; + /// All static element indices in the array type, in the ArrayType's canonical order. + SmallVector indices; + /// The final SSA value at each element index, from either the initializer or a later write. + DenseMap valueByIndex; + /// The type of the final value at each element index. + DenseMap typeByIndex; + /// The write operation that overwrites each element index, used for dominance-like ordering + /// checks. Indices defined only by the array initializer have no entry. + DenseMap writeOpByIndex; + /// Discardable attributes from the write that supplies each element value. These are retained + /// because the element writes are erased after their scalar member-write replacements are made. + DenseMap writeDiscardableAttrsByIndex; + /// Direct writes to the local allocation. + SmallVector writes; + /// Direct reads from the local allocation. + SmallVector reads; + /// Direct writes that store the whole local allocation into a struct member. + SmallVector memberWrites; +}; - LogicalResult matchAndRewrite(CreateArrayOp op, PatternRewriter &rewriter) const override { - Value createResult = op.getResult(); - ArrayType createResultType = dyn_cast(createResult.getType()); - assert(createResultType && "CreateArrayOp must produce ArrayType"); - Type oldResultElemType = createResultType.getElementType(); +/// Scalar member name and type. +using MemberInfo = std::pair; - // Look for WriteArrayOp where the array reference is the result of the CreateArrayOp and the - // element type is different. - Type newResultElemType = nullptr; - for (Operation *user : createResult.getUsers()) { - if (WriteArrayOp writeOp = dyn_cast(user)) { - if (writeOp.getArrRef() != createResult) { - continue; - } - Type writeRValueType = writeOp.getRvalue().getType(); - if (writeRValueType == oldResultElemType) { - continue; - } - if (newResultElemType && newResultElemType != writeRValueType) { - LLVM_DEBUG( - llvm::dbgs() +/// Replacement scalar members for one array-typed member. +struct SplitMemberInfo { + /// Static element indices that were split out of the original array-typed member. + SmallVector indices; + /// Replacement scalar member for each static element index. + DenseMap memberByIndex; +}; + +/// Return true iff `lhs` and `rhs` contain the same array indices. +static bool haveSameIndexSet(ArrayRef lhs, ArrayRef rhs) { + if (lhs.size() != rhs.size()) { + return false; + } + DenseSet rhsSet(rhs.begin(), rhs.end()); + return llvm::all_of(lhs, [&rhsSet](ArrayAttr idx) { return rhsSet.contains(idx); }); +} + +/// Return the first index in `lhs` that is not present in `rhs`. +static ArrayAttr findIndexMissingFrom(ArrayRef lhs, ArrayRef rhs) { + DenseSet rhsSet(rhs.begin(), rhs.end()); + const auto *it = llvm::find_if(lhs, [&rhsSet](ArrayAttr idx) { return !rhsSet.contains(idx); }); + return it == lhs.end() ? ArrayAttr() : *it; +} + +/// Append indices from `source` that are not already present in `target`. +static void appendMissingIndices(SmallVector &target, ArrayRef source) { + DenseSet targetSet(target.begin(), target.end()); + for (ArrayAttr idx : source) { + if (targetSet.insert(idx).second) { + target.push_back(idx); + } + } +} + +/// Replace all uses of `oldValue` without asking MLIR to enforce SSA type equality. +/// +/// This is intentionally narrower than `replaceAllUsesWith()`: the whole purpose of this step is to +/// remove pseudo-homogeneous array values whose original element type no longer describes the +/// concrete value stored at each index. The caller must have already proven that every rewritten +/// use observes the value at a single static index, so replacing that use with the index-specific +/// value is type-correct for the consuming operation after the rewrite. +static inline void replaceAllUsesIgnoringType(Value oldValue, Value newValue) { + for (OpOperand &use : llvm::make_early_inc_range(oldValue.getUses())) { + use.set(newValue); + } +} + +/// Return true iff replacing `readOp` with a value of `replacementType` preserves the read result's +/// type refinement. +static inline bool canReplaceReadResultWithType( + ReadArrayOp readOp, Type replacementType, const ConversionTracker &tracker, const char *patName +) { + Type readResultType = readOp.getResult().getType(); + return readResultType == replacementType || + tracker.isLegalConversion(readResultType, replacementType, patName); +} + +/// Return true iff direct typed consumers of `readOp` accept `replacementType`. +/// +/// Scalarizing a pseudo-homogeneous array can replace a generic read result with the concrete +/// value stored at its static index. Checking the read result alone is insufficient: two +/// different concrete types can both unify with that generic result, while a consumer may require +/// only one of them. Typed consumers with separately declared requirements must be checked before +/// the type-ignoring replacement below. +static LogicalResult canReplaceReadUsersWithType( + ReadArrayOp readOp, Type replacementType, SymbolTableCollection &tables, + const DenseMap *candidateInfoByCreateOp = nullptr +) { + Value result = readOp.getResult(); + for (OpOperand &use : result.getUses()) { + Operation *user = use.getOwner(); + if (MemberWriteOp memberWrite = llvm::dyn_cast(user)) { + if (memberWrite.getVal() != result) { + return failure(); + } + auto memberDef = memberWrite.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + Type memberType = memberDef->get().getType(); + // A parameterized member will be refined by the normal propagation patterns. Only a + // concrete member independently constrains which index-specific value is valid here. + if (isConcreteType(memberType, /*allowStructParams=*/false) && + !typesUnify(replacementType, memberType)) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with a member write" + ); + diag.attachNote(memberWrite.getLoc()) + << "member write requires " << memberType << ", but this read is replaced with " + << replacementType; + return diag; + } + continue; + } + if (CallOp call = llvm::dyn_cast(user)) { + unsigned argIdx = use.getOperandNumber() - call.getArgOperands().getBeginOperandIndex(); + if (argIdx >= call.getArgOperands().size()) { + return failure(); + } + auto callee = call.getCalleeTarget(tables); + if (failed(callee)) { + return failure(); + } + Type paramType = callee->get().getFunctionType().getInput(argIdx); + if (isConcreteType(paramType, /*allowStructParams=*/false) && + !typesUnify(replacementType, paramType, callee->getNamespace())) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with a function call" + ); + diag.attachNote(call.getLoc()) << "call argument requires " << paramType + << ", but this read is replaced with " << replacementType; + return diag; + } + continue; + } + if (WriteArrayOp arrayWrite = llvm::dyn_cast(user)) { + if (arrayWrite.getRvalue() != result) { + return failure(); + } + Type elementType = arrayWrite.getArrRefType().getElementType(); + if (isConcreteType(elementType, /*allowStructParams=*/false) && + !typesUnify(replacementType, elementType)) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with an array write" + ); + diag.attachNote(arrayWrite.getLoc()) + << "array write requires " << elementType << ", but this read is replaced with " + << replacementType; + return diag; + } + continue; + } + if (CreateArrayOp createArray = llvm::dyn_cast(user)) { + if (use.getOperandNumber() >= createArray.getElements().size()) { + return failure(); + } + // A downstream scalarization candidate consumes the replacement while it is still valid, + // then removes the initialized array. Its cached initializer values are refreshed from this + // candidate before any rewrites run, so it is safe to let that supported path proceed. + if (candidateInfoByCreateOp && + candidateInfoByCreateOp->contains(createArray.getOperation())) { + continue; + } + Type elementType = createArray.getType().getElementType(); + // Unlike array.write, array.new requires each initializer to exactly match the result + // array's element type. A generic read may be a valid initializer before scalarization, + // even though its index-specific replacement is not. + if (replacementType != elementType) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with an array initializer" + ); + diag.attachNote(createArray.getLoc()) + << "array initializer requires exactly " << elementType + << ", but this read is replaced with " << replacementType; + return diag; + } + continue; + } + if (UnifiableCastOp castOp = llvm::dyn_cast(user)) { + if (castOp.getInput() != result) { + return failure(); + } + Type castResultType = castOp.getResult().getType(); + // The cast's result may independently require a concrete type. Replacing its generic input + // with an index-specific value that cannot unify with that result would make the cast + // invalid. + if (isConcreteType(castResultType, /*allowStructParams=*/false) && + !typesUnify(replacementType, castResultType)) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with a unifiable cast" + ); + diag.attachNote(castOp.getLoc()) << "unifiable cast result requires " << castResultType + << ", but this read is replaced with " << replacementType; + return diag; + } + continue; + } + if (ReturnOp returnOp = llvm::dyn_cast(user)) { + unsigned resultIdx = use.getOperandNumber(); + FuncDefOp function = returnOp->getParentOfType(); + TypeRange resultTypes = function.getFunctionType().getResults(); + if (resultIdx >= resultTypes.size()) { + return failure(); + } + Type resultType = resultTypes[resultIdx]; + if (isConcreteType(resultType, /*allowStructParams=*/false) && + !typesUnify(replacementType, resultType)) { + InFlightDiagnostic diag = readOp.emitError( + "cannot scalarize heterogeneous array read because its index-specific value type is " + "incompatible with a function return" + ); + diag.attachNote(returnOp.getLoc()) + << "function return requires " << resultType << ", but this read is replaced with " + << replacementType; + return diag; + } + } + } + return success(); +} + +/// Return true if `def` is in the same block as `user` and appears before it. +static inline bool strictlyBefore(Operation *def, Operation *user) { + return def->getBlock() == user->getBlock() && def->isBeforeInBlock(user); +} + +/// Return all direct users of `value`, sorted by their order in the containing block. +static FailureOr> getUsersInBlockOrder(Value value) { + SmallVector users(value.getUsers().begin(), value.getUsers().end()); + if (users.empty()) { + return users; + } + Block *block = users.front()->getBlock(); + if (!llvm::all_of(users, [block](Operation *user) { return user->getBlock() == block; })) { + return failure(); + } + llvm::sort(users, [](Operation *lhs, Operation *rhs) { return lhs->isBeforeInBlock(rhs); }); + return users; +} + +/// Return true iff `createOp` has direct element writes whose order can update stored values. +static inline bool hasDirectArrayWrites(CreateArrayOp createOp) { + Value arrayValue = createOp.getResult(); + return llvm::any_of(arrayValue.getUsers(), [arrayValue](Operation *user) { + auto writeOp = llvm::dyn_cast(user); + return writeOp && writeOp.getArrRef() == arrayValue; + }); +} + +/// Return true iff `createOp` provides an initializer value for every static array element. +static inline bool hasFullInitializer(CreateArrayOp createOp, ArrayRef indices) { + Operation::operand_range elements = createOp.getElements(); + return !elements.empty() && elements.size() == indices.size(); +} + +/// Return true iff all element values are immutable initializer operands. +static inline bool +isInitializerOnlyLocalArray(CreateArrayOp createOp, ArrayRef indices) { + return hasFullInitializer(createOp, indices) && !hasDirectArrayWrites(createOp); +} + +/// Return true iff the array allocation and every initializer value dominate `user`. +static inline bool initializerValuesDominateUse( + CreateArrayOp createOp, Operation *user, const DominanceInfo &domInfo +) { + if (!domInfo.dominates(createOp.getResult(), user)) { + return false; + } + return llvm::all_of(createOp.getElements(), [user, &domInfo](Value element) { + return domInfo.dominates(element, user); + }); +} + +/// Return true if all writes for the scalarized allocation are available before `user`. +/// +/// The rewrite currently handles straight-line local array construction. Requiring every write to +/// be in the same block and before the consuming read/member write keeps the replacement local and +/// avoids changing behavior for arrays updated through control flow. +static inline bool allWritesAvailableAt(const ScalarizedArrayInfo &info, Operation *user) { + return llvm::all_of(info.writeOpByIndex, [user](const auto &entry) { + return strictlyBefore(entry.second, user); + }); +} + +/// Convert array access operands to a static index attribute, if possible. +static inline ArrayAttr getIndexAsAttr(ArrayAccessOpInterface op) { + return op.indexOperandsToAttributeArray(); +} + +static Type +specializeTypeForArrayIndex(Type type, ArrayAttr idx, const ConversionTracker *tracker = nullptr); + +/// Fold an affine-map attribute against the static array index, if the map only depends on that +/// index. +static Attribute +specializeAttrForArrayIndex(Attribute attr, ArrayAttr idx, const ConversionTracker *tracker) { + if (!attr) { + return attr; + } + if (auto mapAttr = llvm::dyn_cast(attr)) { + AffineMap map = mapAttr.getAffineMap(); + if (idx.size() != map.getNumDims() + map.getNumSymbols()) { + return attr; + } + SmallVector operands; + operands.reserve(idx.size()); + for (Attribute idxPart : idx.getValue()) { + auto intAttr = llvm::dyn_cast(idxPart); + if (!intAttr) { + return attr; + } + operands.push_back(intAttr); + } + + SmallVector result; + bool hasPoison = false; + if (failed(map.constantFold(operands, result, &hasPoison)) || hasPoison || result.size() != 1) { + return attr; + } + return result.front(); + } + if (auto typeAttr = llvm::dyn_cast(attr)) { + Type specialized = specializeTypeForArrayIndex(typeAttr.getValue(), idx, tracker); + return specialized == typeAttr.getValue() ? attr : TypeAttr::get(specialized); + } + if (auto arrayAttr = llvm::dyn_cast(attr)) { + MLIRContext *ctx = attr.getContext(); + SmallVector specializedAttrs; + bool changed = false; + for (Attribute nested : arrayAttr.getValue()) { + Attribute specialized = specializeAttrForArrayIndex(nested, idx, tracker); + specializedAttrs.push_back(specialized); + changed |= specialized != nested; + } + return changed ? ArrayAttr::get(ctx, specializedAttrs) : attr; + } + return attr; +} + +/// Specialize affine-map parameters nested in `type` for the given static array index. +static Type +specializeTypeForArrayIndex(Type type, ArrayAttr idx, const ConversionTracker *tracker) { + if (auto structTy = llvm::dyn_cast(type)) { + ArrayAttr params = structTy.getParams(); + if (!params) { + return type; + } + SmallVector specializedParams; + bool changed = false; + for (Attribute param : params.getValue()) { + Attribute specialized = specializeAttrForArrayIndex(param, idx, tracker); + specializedParams.push_back(specialized); + changed |= specialized != param; + } + if (!changed) { + return type; + } + Type specialized = StructType::get( + structTy.getNameRef(), ArrayAttr::get(type.getContext(), specializedParams) + ); + if (tracker) { + if (auto structInstantiation = + tracker->getInstantiation(llvm::cast(specialized))) { + return *structInstantiation; + } + } + return specialized; + } + if (auto arrayTy = llvm::dyn_cast(type)) { + SmallVector specializedDims; + bool dimsChanged = false; + for (Attribute dim : arrayTy.getDimensionSizes()) { + Attribute specialized = specializeAttrForArrayIndex(dim, idx, tracker); + specializedDims.push_back(specialized); + dimsChanged |= specialized != dim; + } + Type specializedElem = specializeTypeForArrayIndex(arrayTy.getElementType(), idx, tracker); + return (dimsChanged || specializedElem != arrayTy.getElementType()) + ? ArrayType::get(specializedElem, specializedDims) + : type; + } + return type; +} + +/// Seed the scalar value/type maps from the explicit elements of a statically-shaped `array.new`. +static LogicalResult seedValuesFromArrayElements( + CreateArrayOp createOp, ArrayRef indices, DenseMap &valueByIndex, + DenseMap *typeByIndex = nullptr, const ConversionTracker *tracker = nullptr +) { + Operation::operand_range elements = createOp.getElements(); + if (elements.empty()) { + return success(); + } + if (elements.size() != indices.size()) { + return failure(); + } + for (auto [idx, value] : llvm::zip_equal(indices, elements)) { + valueByIndex[idx] = value; + if (typeByIndex) { + Type specializedType = + specializeTypeForArrayIndex(createOp.getType().getElementType(), idx, tracker); + bool canUseSpecializedType = + typesUnify(value.getType(), specializedType) || + (tracker && tracker->isLegalConversion( + value.getType(), specializedType, "seedValuesFromArrayElements" + )); + (*typeByIndex)[idx] = canUseSpecializedType ? specializedType : value.getType(); + } + } + return success(); +} + +/// Return true iff the candidate array stores at least two non-unifying element types. +/// +/// Homogeneous arrays should continue through the normal type-propagation path. This step exists +/// specifically for pseudo-homogeneous arrays, such as arrays whose element is a templated struct +/// with an affine-map parameter that becomes a different concrete struct at each unrolled index. +static bool hasMultipleIncompatibleElementTypes(const ScalarizedArrayInfo &info) { + SmallVector previousTypes; + for (ArrayAttr idx : info.indices) { + Type nextType = info.typeByIndex.lookup(idx); + if (!nextType) { + return false; + } + for (Type previousType : previousTypes) { + if (!typesUnify(previousType, nextType)) { + return true; + } + } + previousTypes.push_back(nextType); + } + return false; +} + +static FailureOr getCommonRefinedType(Type lhs, Type rhs, const ConversionTracker &tracker); + +/// Return the common refinement of two matching type parameters, if it can be represented directly. +static FailureOr getCommonRefinedParamAttr( + Attribute lhs, Attribute rhs, const ConversionTracker &tracker, bool unifyDynamicSize = false +) { + assertValidAttrForParamOfType(lhs); + assertValidAttrForParamOfType(rhs); + if (lhs == rhs) { + return lhs; + } + + if (auto lhsAffine = llvm::dyn_cast(lhs)) { + if (auto rhsInt = llvm::dyn_cast(rhs)) { + if (!isDynamic(rhsInt)) { + return rhs; + } + } + } + if (auto rhsAffine = llvm::dyn_cast(rhs)) { + if (auto lhsInt = llvm::dyn_cast(lhs)) { + if (!isDynamic(lhsInt)) { + return lhs; + } + } + } + + if (unifyDynamicSize) { + auto dynamicInt = [](Attribute attr) -> IntegerAttr { + auto intAttr = llvm::dyn_cast(attr); + return intAttr && isDynamic(intAttr) ? intAttr : nullptr; + }; + if (dynamicInt(lhs) && llvm::isa_and_present(rhs)) { + return rhs; + } + if (dynamicInt(rhs) && llvm::isa_and_present(lhs)) { + return lhs; + } + } + + auto lhsTy = llvm::dyn_cast(lhs); + auto rhsTy = llvm::dyn_cast(rhs); + if (lhsTy && rhsTy) { + auto commonType = getCommonRefinedType(lhsTy.getValue(), rhsTy.getValue(), tracker); + if (succeeded(commonType)) { + return TypeAttr::get(*commonType); + } + } + return failure(); +} + +/// Return the common refinement of two unifying types, if this pass can name it safely. +static FailureOr getCommonRefinedType(Type lhs, Type rhs, const ConversionTracker &tracker) { + if (lhs == rhs) { + return lhs; + } + if (tracker.isLegalConversion(lhs, rhs, "getCommonRefinedType")) { + return rhs; + } + if (tracker.isLegalConversion(rhs, lhs, "getCommonRefinedType")) { + return lhs; + } + + if (auto lhsStruct = llvm::dyn_cast(lhs)) { + auto rhsStruct = llvm::dyn_cast(rhs); + if (!rhsStruct) { + return failure(); + } + if (lhsStruct.getNameRef() != rhsStruct.getNameRef()) { + std::optional lhsPreimage = tracker.getPreimage(lhsStruct); + std::optional rhsPreimage = tracker.getPreimage(rhsStruct); + if (!lhsPreimage && !rhsPreimage) { + return failure(); + } + auto commonPreimage = getCommonRefinedType( + lhsPreimage.value_or(lhsStruct), rhsPreimage.value_or(rhsStruct), tracker + ); + if (failed(commonPreimage)) { + return failure(); + } + if (auto commonStruct = llvm::dyn_cast(*commonPreimage)) { + if (std::optional instantiatedCommon = tracker.getInstantiation(commonStruct)) { + return *instantiatedCommon; + } + } + return *commonPreimage; + } + + ArrayAttr lhsParams = lhsStruct.getParams(); + ArrayAttr rhsParams = rhsStruct.getParams(); + ArrayRef emptyParams; + ArrayRef lhsValues = lhsParams ? lhsParams.getValue() : emptyParams; + ArrayRef rhsValues = rhsParams ? rhsParams.getValue() : emptyParams; + if (lhsValues.size() != rhsValues.size()) { + return failure(); + } + + SmallVector commonParams; + for (auto [lhsParam, rhsParam] : llvm::zip_equal(lhsValues, rhsValues)) { + auto commonParam = getCommonRefinedParamAttr(lhsParam, rhsParam, tracker); + if (failed(commonParam)) { + return failure(); + } + commonParams.push_back(*commonParam); + } + StructType commonStruct = + commonParams.empty() + ? StructType::get(lhsStruct.getNameRef()) + : StructType::get( + lhsStruct.getNameRef(), ArrayAttr::get(lhs.getContext(), commonParams) + ); + if (std::optional instantiatedCommon = tracker.getInstantiation(commonStruct)) { + return *instantiatedCommon; + } + return commonStruct; + } + + if (auto lhsArray = llvm::dyn_cast(lhs)) { + auto rhsArray = llvm::dyn_cast(rhs); + if (!rhsArray) { + return failure(); + } + + auto commonElement = + getCommonRefinedType(lhsArray.getElementType(), rhsArray.getElementType(), tracker); + if (failed(commonElement)) { + return failure(); + } + ArrayRef lhsDims = lhsArray.getDimensionSizes(); + ArrayRef rhsDims = rhsArray.getDimensionSizes(); + if (lhsDims.size() != rhsDims.size()) { + return failure(); + } + + SmallVector commonDims; + for (auto [lhsDim, rhsDim] : llvm::zip_equal(lhsDims, rhsDims)) { + auto commonDim = + getCommonRefinedParamAttr(lhsDim, rhsDim, tracker, /*unifyDynamicSize=*/true); + if (failed(commonDim)) { + return failure(); + } + commonDims.push_back(*commonDim); + } + return ArrayType::get(*commonElement, commonDims); + } + + return failure(); +} + +/// Return true iff a value of `sourceType` can be used where `targetType` is requested after +/// scalarization has retained the common refinement. +static bool canUseScalarizedValueAsType( + Type sourceType, Type targetType, const ConversionTracker &tracker, const char *patName +) { + if (sourceType == targetType || tracker.isLegalConversion(sourceType, targetType, patName)) { + return true; + } + auto commonType = getCommonRefinedType(sourceType, targetType, tracker); + return succeeded(commonType) && *commonType == targetType; +} + +/// Return true iff `type` has no symbolic array dimensions or element type. +static bool isFullyConcreteArrayConsumerType(Type type) { + auto arrayType = llvm::dyn_cast(type); + if (!arrayType) { + return isConcreteType(type, /*allowStructParams=*/false); + } + return llvm::all_of(arrayType.getDimensionSizes(), isConcreteAttr<>) && + isFullyConcreteArrayConsumerType(arrayType.getElementType()); +} + +/// Return true iff typed users of `array.new` can accept a refined result type. +/// +/// Retagging an array result updates its SSA value globally. Fully concrete consumers are not +/// retagged by propagation patterns, so they must already request the proposed array type. +static bool canRefineCreateArrayUsersToType( + CreateArrayOp createOp, Type refinedType, SymbolTableCollection &tables +) { + Value result = createOp.getResult(); + for (OpOperand &use : result.getUses()) { + Operation *user = use.getOwner(); + Type requiredType; + if (MemberWriteOp memberWrite = llvm::dyn_cast(user)) { + if (memberWrite.getVal() != result) { + return false; + } + auto memberDef = memberWrite.getMemberDefOp(tables); + if (failed(memberDef)) { + return false; + } + requiredType = memberDef->get().getType(); + } else if (CallOp call = llvm::dyn_cast(user)) { + unsigned argIdx = use.getOperandNumber() - call.getArgOperands().getBeginOperandIndex(); + if (argIdx >= call.getArgOperands().size()) { + return false; + } + auto callee = call.getCalleeTarget(tables); + if (failed(callee)) { + return false; + } + requiredType = callee->get().getFunctionType().getInput(argIdx); + } else if (ReturnOp returnOp = llvm::dyn_cast(user)) { + FuncDefOp function = returnOp->getParentOfType(); + unsigned resultIdx = use.getOperandNumber(); + TypeRange resultTypes = function.getFunctionType().getResults(); + if (resultIdx >= resultTypes.size()) { + return false; + } + requiredType = resultTypes[resultIdx]; + } else if (UnifiableCastOp castOp = llvm::dyn_cast(user)) { + if (castOp.getInput() != result) { + return false; + } + requiredType = castOp.getResult().getType(); + } else { + continue; + } + // Array element specializations with distinct concrete affine arguments can still unify + // before template instantiation. This array result is retagged in place, though, and no + // propagation pattern retargets these users, so a fully concrete requirement must match + // exactly. + if (isFullyConcreteArrayConsumerType(requiredType) && refinedType != requiredType) { + return false; + } + } + return true; +} + +/// Return true iff typed consumers of a member read can use its refined result type. +/// +/// A generic read result may unify both with the proposed member type and with a different +/// concrete type required by a consumer. Updating the read in that situation would make the +/// consuming operation invalid, so the member definition must remain generic. +static bool canRefineMemberReadUsersToType( + MemberReadOp readOp, Type refinedType, SymbolTableCollection &tables, + const ConversionTracker &tracker +) { + Value result = readOp.getVal(); + for (OpOperand &use : result.getUses()) { + Operation *user = use.getOwner(); + if (MemberWriteOp memberWrite = llvm::dyn_cast(user)) { + if (memberWrite.getVal() != result) { + return false; + } + auto memberDef = memberWrite.getMemberDefOp(tables); + if (failed(memberDef)) { + return false; + } + Type memberType = memberDef->get().getType(); + if (isConcreteType(memberType, /*allowStructParams=*/false) && + !canUseScalarizedValueAsType( + refinedType, memberType, tracker, "UpdateMemberDefTypeFromWrite" + )) { + return false; + } + continue; + } + if (CallOp call = llvm::dyn_cast(user)) { + unsigned argIdx = use.getOperandNumber() - call.getArgOperands().getBeginOperandIndex(); + if (argIdx >= call.getArgOperands().size()) { + return false; + } + auto callee = call.getCalleeTarget(tables); + if (failed(callee)) { + return false; + } + Type paramType = callee->get().getFunctionType().getInput(argIdx); + if (isConcreteType(paramType, /*allowStructParams=*/false) && + !canUseScalarizedValueAsType( + refinedType, paramType, tracker, "UpdateMemberDefTypeFromWrite" + )) { + return false; + } + continue; + } + if (CreateArrayOp createArray = llvm::dyn_cast(user)) { + if (use.getOperandNumber() >= createArray.getElements().size()) { + return false; + } + // array.new requires every initializer to exactly match its declared element type. Allow + // this refinement only when the array can follow it: every other initializer already has + // the refined type and the array element type can be refined accordingly. + bool canUse = canUseScalarizedValueAsType( + createArray.getType().getElementType(), refinedType, tracker, + "UpdateMemberDefTypeFromWrite" + ); + if (!canUse) { + return false; + } + for (OpOperand &initializer : createArray->getOpOperands()) { + if (&initializer != &use && initializer.get().getType() != refinedType) { + return false; + } + } + ArrayType refinedArrayType = createArray.getType().cloneWith(refinedType); + if (!canRefineCreateArrayUsersToType(createArray, refinedArrayType, tables)) { + return false; + } + continue; + } + if (WriteArrayOp arrayWrite = llvm::dyn_cast(user)) { + if (arrayWrite.getRvalue() != result) { + return false; + } + Type elementType = arrayWrite.getArrRefType().getElementType(); + // array.write's rvalue is not independently converted by the member refinement. Its + // destination must therefore already require exactly the proposed read type. In + // particular, do not use the normal unification-based conversion check here: symbolic + // Cell<0> and Cell<1> can appear compatible before their template instantiations are + // materialized, but become distinct concrete element types later in flattening. + if (refinedType != elementType) { + return false; + } + continue; + } + } + return true; +} + +/// Merge one candidate scalar type into the split type map, keeping the most concrete refinement. +static LogicalResult mergeSplitCandidateType( + MemberDefOp member, ArrayAttr idx, Type candidateType, Operation *candidateOp, + DenseMap &splitTypes, DenseMap *splitTypeOps, + const ConversionTracker &tracker +) { + auto existing = splitTypes.find(idx); + if (existing == splitTypes.end()) { + splitTypes[idx] = candidateType; + if (splitTypeOps) { + (*splitTypeOps)[idx] = candidateOp; + } + return success(); + } + + Type existingType = existing->second; + if (existingType == candidateType) { + return success(); + } + if (tracker.isLegalConversion(existingType, candidateType, "mergeSplitCandidateType")) { + existing->second = candidateType; + if (splitTypeOps) { + (*splitTypeOps)[idx] = candidateOp; + } + return success(); + } + if (tracker.isLegalConversion(candidateType, existingType, "mergeSplitCandidateType") || + existingType == candidateType) { + return success(); + } + + auto commonType = getCommonRefinedType(existingType, candidateType, tracker); + if (succeeded(commonType)) { + existing->second = *commonType; + if (splitTypeOps) { + (*splitTypeOps)[idx] = candidateOp; + } + return success(); + } + + InFlightDiagnostic diag = member.emitError( + "cannot split heterogeneous array member because candidate writes require incompatible " + "scalar member types" + ); + if (splitTypeOps) { + diag.attachNote(splitTypeOps->lookup(idx)->getLoc()) + << "candidate writes index " << idx << " with type " << existingType; + diag.attachNote(candidateOp->getLoc()) + << "conflicting candidate writes the same index with type " << candidateType; + } + return diag; +} + +/// Collect scalarization information for `op` if it is a safe heterogeneous-array candidate. +/// +/// A candidate must: +/// - have a static shape and a real element type, +/// - have only direct reads, direct writes, and whole-array struct member writes as users, +/// - use only static array indices, +/// - have an initializer or exactly one write for every static element index, +/// - have all reads/member writes after every explicit element write, and +/// - store multiple incompatible element types. +static FailureOr +getScalarizedArrayInfo(CreateArrayOp op, const ConversionTracker &tracker) { + ArrayType arrTy = op.getType(); + if (!arrTy.hasStaticShape() || llvm::isa(arrTy.getElementType())) { + return failure(); + } + + std::optional> maybeIndices = arrTy.getSubelementIndices(); + if (!maybeIndices) { + return failure(); + } + + ScalarizedArrayInfo info; + info.createOp = op; + info.indices = std::move(*maybeIndices); + auto seedRes = + seedValuesFromArrayElements(op, info.indices, info.valueByIndex, &info.typeByIndex, &tracker); + if (failed(seedRes)) { + return failure(); + } + Value arrayValue = op.getResult(); + + for (Operation *user : arrayValue.getUsers()) { + if (auto writeOp = llvm::dyn_cast(user)) { + if (writeOp.getArrRef() != arrayValue) { + return failure(); + } + ArrayAttr idx = getIndexAsAttr(writeOp); + if (!idx) { + return failure(); + } + if (info.writeOpByIndex.contains(idx)) { + return failure(); + } + info.valueByIndex[idx] = writeOp.getRvalue(); + info.typeByIndex[idx] = writeOp.getRvalue().getType(); + info.writeOpByIndex[idx] = writeOp.getOperation(); + info.writeDiscardableAttrsByIndex[idx] = writeOp->getDiscardableAttrDictionary(); + info.writes.push_back(writeOp); + continue; + } + if (auto readOp = llvm::dyn_cast(user)) { + if (readOp.getArrRef() != arrayValue) { + return failure(); + } + ArrayAttr idx = getIndexAsAttr(readOp); + if (!idx) { + return failure(); + } + info.reads.push_back(readOp); + continue; + } + if (auto memberWriteOp = llvm::dyn_cast(user)) { + if (memberWriteOp.getVal() != arrayValue) { + return failure(); + } + info.memberWrites.push_back(memberWriteOp); + continue; + } + return failure(); + } + + for (ArrayAttr idx : info.indices) { + if (!info.valueByIndex.contains(idx)) { + return failure(); + } + } + for (ReadArrayOp readOp : info.reads) { + ArrayAttr idx = getIndexAsAttr(readOp); + if (!info.valueByIndex.contains(idx) || !allWritesAvailableAt(info, readOp)) { + return failure(); + } + if (!canReplaceReadResultWithType( + readOp, info.typeByIndex.lookup(idx), tracker, "getScalarizedArrayInfo" + )) { + return failure(); + } + } + for (MemberWriteOp memberWriteOp : info.memberWrites) { + if (!allWritesAvailableAt(info, memberWriteOp)) { + return failure(); + } + } + if (!hasMultipleIncompatibleElementTypes(info)) { + return failure(); + } + return info; +} + +/// Create scalar replacement members for `member`, or return the replacements already created. +/// +/// The replacement members preserve the original member's public/signal/column and discardable +/// metadata, and rely on the containing struct's symbol table to make each generated name unique. +static SplitMemberInfo &getOrCreateSplitMemberInfo( + MemberDefOp member, ScalarizedArrayInfo &arrayInfo, + const DenseMap> &splitTypesByMember, + const DenseMap> &splitIndicesByMember, + DenseMap &splitMembers, SymbolTableCollection &tables, + PatternRewriter &rewriter, const ConversionTracker &tracker +) { + auto existing = splitMembers.find(member); + if (existing != splitMembers.end()) { + return existing->second; + } + + SplitMemberInfo &splitInfo = splitMembers[member]; + if (auto indicesIt = splitIndicesByMember.find(member); indicesIt != splitIndicesByMember.end()) { + splitInfo.indices = indicesIt->second; + } else { + splitInfo.indices = arrayInfo.indices; + } + + StructDefOp parentStruct = getParentOfType(member); + assert(parentStruct && "MemberDefOp parent is always StructDefOp"); + SymbolTable &structSymbols = tables.getSymbolTable(parentStruct); + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(member); + auto splitTypesIt = splitTypesByMember.find(member); + for (ArrayAttr idx : splitInfo.indices) { + Type scalarType = arrayInfo.typeByIndex.lookup(idx); + if (splitTypesIt != splitTypesByMember.end()) { + if (Type refinedType = splitTypesIt->second.lookup(idx)) { + scalarType = refinedType; + } + } + if (!scalarType) { + scalarType = + specializeTypeForArrayIndex(arrayInfo.createOp.getType().getElementType(), idx, &tracker); + } + MemberDefOp newMember = rewriter.create( + member.getLoc(), member.getSymNameAttr(), scalarType, member.getSignal(), member.getColumn() + ); + newMember.setPublicAttr(member.hasPublicAttr()); + newMember->setDiscardableAttrs(member->getDiscardableAttrDictionary()); + StringAttr actualName = structSymbols.insert(newMember); + splitInfo.memberByIndex[idx] = std::make_pair(actualName, scalarType); + } + return splitInfo; +} + +/// Value and type cached for one scalarized array index. +struct CachedScalarizedValue { + /// SSA value currently known to provide this scalarized index. + Value value; + /// Type that the scalarized index must use when splitting member writes. + Type type; +}; + +/// Return the cached scalar value for `value` when it is a static read from another candidate. +static std::optional getCandidateReadReplacement( + Value value, const DenseMap *candidateInfoByCreateOp +) { + if (!candidateInfoByCreateOp) { + return std::nullopt; + } + auto readOp = value.getDefiningOp(); + if (!readOp) { + return std::nullopt; + } + auto createOp = readOp.getArrRef().getDefiningOp(); + if (!createOp) { + return std::nullopt; + } + auto infoIt = candidateInfoByCreateOp->find(createOp.getOperation()); + if (infoIt == candidateInfoByCreateOp->end()) { + return std::nullopt; + } + ArrayAttr idx = getIndexAsAttr(readOp); + if (!idx) { + return std::nullopt; + } + + ScalarizedArrayInfo *info = infoIt->second; + Value replacement = info->valueByIndex.lookup(idx); + Type replacementType = info->typeByIndex.lookup(idx); + if (!replacement || !replacementType) { + return std::nullopt; + } + return CachedScalarizedValue {replacement, replacementType}; +} + +/// Refresh cached element values from their still-live array operands. +/// +/// Candidates are collected before any rewrites run, so one candidate can cache a read result from +/// another candidate. Rewriting the upstream candidate updates the downstream write operands and +/// initializer operands and erases the reads; re-reading the operands here keeps this candidate +/// from using dangling values. +static LogicalResult refreshValuesFromArrayOperands( + ScalarizedArrayInfo &info, const ConversionTracker &tracker, + const DenseMap *candidateInfoByCreateOp = nullptr, + bool *changed = nullptr +) { + DenseMap oldValueByIndex; + DenseMap oldTypeByIndex; + if (changed != nullptr) { + oldValueByIndex = info.valueByIndex; + oldTypeByIndex = info.typeByIndex; + } + + auto setCachedValue = [&](ArrayAttr idx, Value value, Type fallbackType) { + if (std::optional replacement = + getCandidateReadReplacement(value, candidateInfoByCreateOp)) { + value = replacement->value; + fallbackType = replacement->type; + } + info.valueByIndex[idx] = value; + info.typeByIndex[idx] = fallbackType; + }; + + Operation::operand_range elements = info.createOp.getElements(); + if (!elements.empty() && elements.size() != info.indices.size()) { + return failure(); + } + if (!elements.empty()) { + for (auto [idx, value] : llvm::zip_equal(info.indices, elements)) { + Type specializedType = + specializeTypeForArrayIndex(info.createOp.getType().getElementType(), idx, &tracker); + bool canUseSpecializedType = + typesUnify(value.getType(), specializedType) || + tracker.isLegalConversion( + value.getType(), specializedType, "refreshValuesFromArrayOperands" + ); + setCachedValue(idx, value, canUseSpecializedType ? specializedType : value.getType()); + } + } + + for (WriteArrayOp writeOp : info.writes) { + ArrayAttr idx = getIndexAsAttr(writeOp); + if (!idx || !info.valueByIndex.contains(idx)) { + return failure(); + } + setCachedValue(idx, writeOp.getRvalue(), writeOp.getRvalue().getType()); + } + if (changed != nullptr) { + for (ArrayAttr idx : info.indices) { + if (oldValueByIndex.lookup(idx) != info.valueByIndex.lookup(idx) || + oldTypeByIndex.lookup(idx) != info.typeByIndex.lookup(idx)) { + *changed = true; + break; + } + } + } + return success(); +} + +/// Refresh all collected candidates through candidate-to-candidate static reads. +static LogicalResult refreshValuesFromDependentCandidates( + MutableArrayRef arraysToScalarize, + DenseMap &candidateInfoByCreateOp, + const ConversionTracker &tracker +) { + candidateInfoByCreateOp.clear(); + for (ScalarizedArrayInfo &info : arraysToScalarize) { + candidateInfoByCreateOp[info.createOp.getOperation()] = &info; + } + + for (unsigned iteration = 0; iteration <= arraysToScalarize.size(); ++iteration) { + bool changed = false; + for (ScalarizedArrayInfo &info : arraysToScalarize) { + auto res = refreshValuesFromArrayOperands(info, tracker, &candidateInfoByCreateOp, &changed); + if (failed(res)) { + return failure(); + } + } + if (!changed) { + return success(); + } + } + return failure(); +} + +/// Verify that refreshing dependent candidates did not invalidate a direct read replacement. +/// +/// A candidate can cache a generic read from another candidate as the value for one of its +/// indices. Refreshing the cache replaces that generic value and type with the upstream +/// candidate's concrete scalar value. Recheck all local reads before any candidates are rewritten +/// so an incompatible replacement cannot fail after earlier candidates have been erased. +static LogicalResult verifyRefreshedCandidateReads( + ArrayRef arraysToScalarize, SymbolTableCollection &tables, + const ConversionTracker &tracker, + const DenseMap &candidateInfoByCreateOp +) { + for (const ScalarizedArrayInfo &info : arraysToScalarize) { + for (ReadArrayOp readOp : info.reads) { + ArrayAttr idx = getIndexAsAttr(readOp); + if (!idx) { + return failure(); + } + bool canReplaceResult = canReplaceReadResultWithType( + readOp, info.typeByIndex.lookup(idx), tracker, "verifyRefreshedCandidateReads" + ); + if (!canReplaceResult) { + return failure(); + } + auto canReplaceUsers = canReplaceReadUsersWithType( + readOp, info.typeByIndex.lookup(idx), tables, &candidateInfoByCreateOp + ); + if (failed(canReplaceUsers)) { + return failure(); + } + } + } + return success(); +} + +/// Collect the common specialized type required for each shared nondeterministic initializer. +/// +/// An inline initializer can store the same generic witness at multiple array indices. Replacing +/// each read or split-member write with a separate typed witness would lose that sharing, so all +/// observed indices using one source witness must agree on a common refinement. In particular, +/// include merged split-member types: another candidate can refine a member even when this +/// candidate's local read remains generic. +static LogicalResult collectSharedNondetSpecializationTypes( + const ScalarizedArrayInfo &info, const ConversionTracker &tracker, + const DenseMap> &splitTypesByMember, + SymbolTableCollection &tables, DenseMap &specializedTypeByNondet +) { + DenseMap> firstUseByNondet; + auto collectType = [&](Value value, Type type, ArrayAttr idx, Location loc) -> LogicalResult { + if (!value || !type || value.getType() == type || !value.getDefiningOp()) { + return success(); + } + auto existing = specializedTypeByNondet.find(value); + if (existing == specializedTypeByNondet.end()) { + specializedTypeByNondet.try_emplace(value, type); + firstUseByNondet.try_emplace(value, std::make_pair(idx, loc)); + return success(); + } + FailureOr commonType = getCommonRefinedType(existing->second, type, tracker); + if (failed(commonType)) { + InFlightDiagnostic diag = info.createOp->emitError( + "cannot scalarize array because a shared nondeterministic initializer requires " + "incompatible specialized types" + ); + auto firstUse = firstUseByNondet.find(value); + assert(firstUse != firstUseByNondet.end() && "first use must be recorded"); + diag.attachNote(firstUse->second.second) + << "array index " << firstUse->second.first << " is read with specialized type " + << existing->second; + diag.attachNote(loc) << "array index " << idx << " is read with specialized type " << type; + return diag; + } + existing->second = *commonType; + return success(); + }; + for (ReadArrayOp readOp : info.reads) { + ArrayAttr idx = getIndexAsAttr(readOp); + auto res = collectType( + info.valueByIndex.lookup(idx), info.typeByIndex.lookup(idx), idx, readOp.getLoc() + ); + if (failed(res)) { + return failure(); + } + } + for (MemberWriteOp memberWriteOp : info.memberWrites) { + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + for (const auto &[idx, type] : splitTypesByMember.lookup(memberDef->get())) { + if (failed(collectType(info.valueByIndex.lookup(idx), type, idx, memberWriteOp.getLoc()))) { + return failure(); + } + } + } + return success(); +} + +/// Verify that each shared witness can be specialized without changing external observations. +static inline bool isInternalSharedNondetUse(const ScalarizedArrayInfo &info, Operation *owner) { + return owner == info.createOp || llvm::any_of(info.writes, [owner](WriteArrayOp writeOp) { + return owner == writeOp.getOperation(); + }); +} + +static LogicalResult verifySharedNondetSpecializationExternalUses( + const ScalarizedArrayInfo &info, const DenseMap &specializedTypeByNondet, + SymbolTableCollection &tables +) { + for (const auto &[source, type] : specializedTypeByNondet) { + for (OpOperand &use : source.getUses()) { + if (isInternalSharedNondetUse(info, use.getOwner())) { + continue; + } + auto call = llvm::dyn_cast(use.getOwner()); + if (!call || !call.calleeIsStructConstrain() || use.getOperandNumber() != 0 || + !llvm::isa(type)) { + InFlightDiagnostic diag = info.createOp->emitError( + "cannot scalarize array because a generic nondeterministic initializer has an " + "unsupported external use" + ); + diag.attachNote(use.getOwner()->getLoc()) << "source witness is also used here"; + return diag; + } + + StructType structType = llvm::cast(type); + SymbolRefAttr callee = + appendLeaf(structType.getNameRef(), call.getCalleeAttr().getLeafReference()); + FailureOr> target = + lookupTopLevelSymbol(tables, callee, call); + if (failed(target)) { + return failure(); + } + + SmallVector argTypes(call.getArgOperands().getTypes()); + argTypes.front() = type; + // Match the namespace-aware compatibility check performed by the CallOp verifier. The + // retargeted callee can accept a still-generic companion argument even when it is not + // exactly equal to the newly instantiated parameter type. + if (!typeListsUnify(argTypes, target->get().getArgumentTypes(), target->getNamespace())) { + InFlightDiagnostic diag = info.createOp->emitError( + "cannot scalarize array because specializing a generic nondeterministic initializer " + "would make an external constraint call incompatible with its retargeted callee" + ); + diag.attachNote(call.getLoc()) + << "retargeted constraint expects argument types " << target->get().getArgumentTypes(); + return diag; + } + } + } + return success(); +} + +/// Materialize each refined witness once and redirect compatible external observations to it. +/// +/// A static array read determines the concrete type of its source witness. Struct constraint calls +/// observing that source can use the same concrete witness after their callee is redirected to the +/// corresponding instantiated struct definition. +static LogicalResult materializeSharedNondetSpecializations( + const ScalarizedArrayInfo &info, const DenseMap &specializedTypeByNondet, + SymbolTableCollection &tables, PatternRewriter &rewriter, + DenseMap &specializedNondetBySource +) { + if (failed(verifySharedNondetSpecializationExternalUses(info, specializedTypeByNondet, tables))) { + return failure(); + } + + DenseMap> externalCallsBySource; + for (const auto &[source, _] : specializedTypeByNondet) { + for (OpOperand &use : source.getUses()) { + if (isInternalSharedNondetUse(info, use.getOwner())) { + continue; + } + externalCallsBySource[source].push_back(llvm::cast(use.getOwner())); + } + } + + for (const auto &[source, type] : specializedTypeByNondet) { + NonDetOp sourceNondet = source.getDefiningOp(); + assert(sourceNondet && "specialization map only contains nondeterministic values"); + OpBuilder::InsertionGuard guard(rewriter); + // The external constraint can precede the array initializer. The source witness dominates + // every one of its original uses, so materializing directly after it preserves SSA dominance + // for both those external calls and the scalarized array reads. + rewriter.setInsertionPointAfter(sourceNondet); + NonDetOp specializedNondet = rewriter.create(sourceNondet.getLoc(), type); + specializedNondet->setDiscardableAttrs(sourceNondet->getDiscardableAttrDictionary()); + Value specializedValue = specializedNondet.getResult(); + specializedNondetBySource[source] = specializedValue; + + for (CallOp call : externalCallsBySource.lookup(source)) { + replaceStructConstraintCallWithSpecializedWitness(rewriter, call, specializedValue); + } + } + return success(); +} + +/// Return the value that should replace `readOp` after scalarization. +/// +/// The cached scalar value can be less concrete than `info.typeByIndex` for inline initializers, +/// so materialize a cast when a concrete consumer needs the refined per-index type. +static FailureOr buildReadReplacementValue( + ReadArrayOp readOp, const ScalarizedArrayInfo &info, PatternRewriter &rewriter, + const ConversionTracker &tracker, const DenseMap &specializedTypeByNondet, + DenseMap &specializedNondetBySource +) { + ArrayAttr idx = getIndexAsAttr(readOp); + if (!idx) { + return failure(); + } + Value replacementValue = info.valueByIndex.lookup(idx); + Type replacementType = info.typeByIndex.lookup(idx); + if (!replacementValue || !replacementType) { + return failure(); + } + if (!canUseScalarizedValueAsType( + replacementValue.getType(), replacementType, tracker, "buildReadReplacementValue" + )) { + return failure(); + } + if (replacementValue.getDefiningOp()) { + // A generic read can require the shared witness refinement solely because another candidate + // contributes a concrete type for the same split-member index. + if (Value cachedNondet = specializedNondetBySource.lookup(replacementValue)) { + return cachedNondet; + } + } + if (replacementValue.getType() == replacementType) { + return replacementValue; + } + + OpBuilder::InsertionGuard guard(rewriter); + if (NonDetOp sourceNondet = replacementValue.getDefiningOp()) { + // Preserve sharing between every index initialized from the same generic witness. + if (Value cachedNondet = specializedNondetBySource.lookup(replacementValue)) { + return cachedNondet; + } + Type specializedType = specializedTypeByNondet.lookup(replacementValue); + if (!specializedType) { + return failure(); + } + // The source witness dominates every original use, including external constraints that can + // precede the array initializer. Materialize directly after it rather than at the first read. + rewriter.setInsertionPointAfter(sourceNondet); + NonDetOp specializedNondet = rewriter.create(readOp.getLoc(), specializedType); + specializedNondet->setDiscardableAttrs(sourceNondet->getDiscardableAttrDictionary()); + specializedNondetBySource[replacementValue] = specializedNondet; + return specializedNondet.getResult(); + } + rewriter.setInsertionPoint(readOp); + if (!typesUnify(replacementValue.getType(), replacementType)) { + return failure(); + } + return rewriter.create(readOp.getLoc(), replacementType, replacementValue) + .getResult(); +} + +/// Emit one scalar `struct.writem` per split index for a whole-array member write. +template +static LogicalResult emitScalarMemberWrites( + MemberWriteOp memberWriteOp, const SplitMemberInfo &splitInfo, + const DenseMap &writeDiscardableAttrsByIndex, + PatternRewriter &rewriter, GetScalarValueFn getScalarValue +) { + rewriter.setInsertionPoint(memberWriteOp); + DictionaryAttr discardableAttrs = memberWriteOp->getDiscardableAttrDictionary(); + for (ArrayAttr idx : splitInfo.indices) { + MemberInfo memberInfo = splitInfo.memberByIndex.lookup(idx); + if (!memberInfo.first) { + return failure(); + } + FailureOr scalarValue = getScalarValue(idx, memberInfo.second); + if (failed(scalarValue)) { + return failure(); + } + MemberWriteOp scalarWrite = rewriter.create( + memberWriteOp.getLoc(), memberWriteOp.getComponent(), + FlatSymbolRefAttr::get(memberInfo.first), *scalarValue + ); + scalarWrite->setDiscardableAttrs(discardableAttrs); + if (DictionaryAttr writeAttrs = writeDiscardableAttrsByIndex.lookup(idx)) { + for (NamedAttribute attr : writeAttrs) { + // The element write is the effective replacement, so preserve its provenance if the + // whole-array and element writes use the same discardable metadata key. + scalarWrite->setDiscardableAttr(attr.getName(), attr.getValue()); + } + } + } + return success(); +} + +/// Return the value for `idx` at `type`, creating and caching a nondeterministic value if the +/// local array element has not been written. +/// +/// The caller must have already rejected initialized expandable arrays that store the same SSA +/// value into indices requiring incompatible split-member types. +static FailureOr getOrCreateScalarizedLocalArrayValue( + DenseMap &valueByIndex, DenseSet &generatedMaterializations, + ArrayAttr idx, Type type, Location loc, PatternRewriter &rewriter, + const ConversionTracker &tracker, Operation *materializationPoint = nullptr +) { + auto createNonDet = [&](Type nondetType, NonDetOp sourceNondet = nullptr) { + OpBuilder::InsertionGuard guard(rewriter); + if (materializationPoint) { + rewriter.setInsertionPointAfter(materializationPoint); + } + NonDetOp nondet = rewriter.create(loc, nondetType); + if (sourceNondet) { + nondet->setDiscardableAttrs(sourceNondet->getDiscardableAttrDictionary()); + } + return nondet.getResult(); + }; + Value scalarValue = valueByIndex.lookup(idx); + if (!scalarValue) { + scalarValue = createNonDet(type); + valueByIndex[idx] = scalarValue; + generatedMaterializations.insert(scalarValue); + } + if (scalarValue.getType() == type) { + return scalarValue; + } + FailureOr commonType = getCommonRefinedType(scalarValue.getType(), type, tracker); + if (failed(commonType)) { + return failure(); + } + if (*commonType == scalarValue.getType()) { + return scalarValue; + } + if (generatedMaterializations.contains(scalarValue)) { + NonDetOp nondetOp = scalarValue.getDefiningOp(); + if (!nondetOp) { + return failure(); + } + rewriter.modifyOpInPlace(nondetOp, [&scalarValue, &commonType]() { + scalarValue.setType(*commonType); + }); + return scalarValue; + } + if (NonDetOp sourceNondet = scalarValue.getDefiningOp()) { + scalarValue = createNonDet(*commonType, sourceNondet); + valueByIndex[idx] = scalarValue; + generatedMaterializations.insert(scalarValue); + } + return scalarValue; +} + +/// Rewrite one local heterogeneous array allocation into its index-specific scalar values. +/// +/// Direct array reads are replaced with the value written at the requested static index. +/// Whole-array member writes are expanded into one scalar member write per index, creating +/// replacement members as needed. Once all consumers are rewritten, the original array writes and +/// allocation are erased. +static LogicalResult rewriteLocalArray( + ScalarizedArrayInfo &info, DenseMap &splitMembers, + const DenseMap> &splitTypesByMember, + const DenseMap> &splitIndicesByMember, + SymbolTableCollection &tables, PatternRewriter &rewriter, const ConversionTracker &tracker, + const DenseMap *candidateInfoByCreateOp +) { + if (failed(refreshValuesFromArrayOperands(info, tracker, candidateInfoByCreateOp))) { + return failure(); + } + + DenseMap specializedTypeByNondet; + if (failed(collectSharedNondetSpecializationTypes( + info, tracker, splitTypesByMember, tables, specializedTypeByNondet + ))) { + return failure(); + } + DenseMap specializedNondetBySource; + if (failed(materializeSharedNondetSpecializations( + info, specializedTypeByNondet, tables, rewriter, specializedNondetBySource + ))) { + return failure(); + } + DenseSet sourcesUsedByReads; + for (ReadArrayOp readOp : info.reads) { + sourcesUsedByReads.insert(info.valueByIndex.lookup(getIndexAsAttr(readOp))); + } + for (ReadArrayOp readOp : llvm::make_early_inc_range(info.reads)) { + ArrayAttr idx = getIndexAsAttr(readOp); + auto canReplaceResult = canReplaceReadResultWithType( + readOp, info.typeByIndex.lookup(idx), tracker, "rewriteLocalArray" + ); + if (!canReplaceResult) { + return failure(); + } + auto canReplaceUsers = canReplaceReadUsersWithType( + readOp, info.typeByIndex.lookup(idx), tables, candidateInfoByCreateOp + ); + if (failed(canReplaceUsers)) { + return failure(); + } + FailureOr replacementValue = buildReadReplacementValue( + readOp, info, rewriter, tracker, specializedTypeByNondet, specializedNondetBySource + ); + if (failed(replacementValue)) { + return failure(); + } + replaceAllUsesIgnoringType(readOp.getResult(), *replacementValue); + rewriter.eraseOp(readOp); + } + + for (MemberWriteOp memberWriteOp : llvm::make_early_inc_range(info.memberWrites)) { + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + SplitMemberInfo &splitInfo = getOrCreateSplitMemberInfo( + memberDef->get(), info, splitTypesByMember, splitIndicesByMember, splitMembers, tables, + rewriter, tracker + ); + + DenseSet generatedMaterializations; + Location memberWriteLoc = memberWriteOp.getLoc(); + auto getScalarValue = [&](ArrayAttr idx, Type type) -> FailureOr { + if (Value scalarValue = info.valueByIndex.lookup(idx)) { + // Reads of a generic initializer materialize one refined witness. Reuse it for the + // corresponding split-member write so both uses retain the original array equality. + if (Value specializedNondet = specializedNondetBySource.lookup(scalarValue)) { + return specializedNondet; + } + return scalarValue; + } + return getOrCreateScalarizedLocalArrayValue( + info.valueByIndex, generatedMaterializations, idx, type, memberWriteLoc, rewriter, tracker + ); + }; + auto emitRes = emitScalarMemberWrites( + memberWriteOp, splitInfo, info.writeDiscardableAttrsByIndex, rewriter, getScalarValue + ); + if (failed(emitRes)) { + return failure(); + } + rewriter.eraseOp(memberWriteOp); + } + + for (WriteArrayOp writeOp : llvm::make_early_inc_range(info.writes)) { + rewriter.eraseOp(writeOp); + } + if (info.createOp.getResult().use_empty()) { + rewriter.eraseOp(info.createOp); + } + for (const auto &[source, _] : specializedTypeByNondet) { + if (!sourcesUsedByReads.contains(source) && source.use_empty()) { + rewriter.eraseOp(source.getDefiningOp()); + } + } + return success(); +} + +/// Rewrite reads from array-typed members that were split by `rewriteLocalArray()`. +/// +/// The only supported use of the original whole-array member read is a static `array.read`. +/// Supporting arbitrary uses would require reconstructing a pseudo-homogeneous array value, which +/// is exactly the invalid representation this step removes. +static LogicalResult rewriteSplitMemberReads( + ModuleOp modOp, DenseMap &splitMembers, + SymbolTableCollection &tables, PatternRewriter &rewriter, const ConversionTracker &tracker +) { + for (MemberReadOp memberReadOp : walkCollect(modOp)) { + auto memberDef = memberReadOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + auto splitIt = splitMembers.find(memberDef->get()); + if (splitIt == splitMembers.end()) { + continue; + } + + SmallVector arrayReads; + for (Operation *user : memberReadOp.getResult().getUsers()) { + auto readOp = llvm::dyn_cast(user); + if (!readOp || readOp.getArrRef() != memberReadOp.getResult() || !getIndexAsAttr(readOp)) { + return failure(); + } + arrayReads.push_back(readOp); + } + + ValueRange mapOperands; + std::optional numDims; + if (!memberReadOp.getMapOperands().empty()) { + mapOperands = memberReadOp.getMapOperands().front(); + numDims = memberReadOp.getNumDimsPerMap().front(); + } + + // A scalar read can be shared only when the source reads have identical + // discardable attributes. These attributes are preserved on the replacement + // operation, so sharing reads with different dictionaries would lose the + // metadata from every read except the last one processed. + DenseMap>> scalarValuesByIndex; + DenseMap scalarValueByRead; + rewriter.setInsertionPoint(memberReadOp); + DictionaryAttr discardableAttrs = memberReadOp->getDiscardableAttrDictionary(); + for (ReadArrayOp readOp : arrayReads) { + ArrayAttr idx = getIndexAsAttr(readOp); + MemberInfo memberInfo = splitIt->second.memberByIndex.lookup(idx); + if (!memberInfo.first) { + return failure(); + } + if (!canReplaceReadResultWithType( + readOp, memberInfo.second, tracker, "rewriteSplitMemberReads" + )) { + return failure(); + } + if (failed(canReplaceReadUsersWithType(readOp, memberInfo.second, tables))) { + return failure(); + } + MemberReadOp scalarRead; + DictionaryAttr readAttrs = readOp->getDiscardableAttrDictionary(); + auto &cachedReads = scalarValuesByIndex[idx]; + auto *cachedIt = llvm::find_if(cachedReads, [&readAttrs](const auto &cachedRead) { + return cachedRead.first == readAttrs; + }); + if (cachedIt != cachedReads.end()) { + scalarRead = llvm::cast(cachedIt->second.getDefiningOp()); + } else { + scalarRead = rewriter.create( + memberReadOp.getLoc(), memberInfo.second, memberReadOp.getComponent(), memberInfo.first, + memberReadOp.getTableOffset().value_or(Attribute {}), mapOperands, numDims + ); + scalarRead->setDiscardableAttrs(discardableAttrs); + // Preserve provenance and other discardable metadata from the static array read that this + // scalar member read replaces. Per-read metadata takes precedence over metadata inherited + // from the whole-array member read. + for (NamedAttribute attr : readAttrs.getValue()) { + scalarRead->setDiscardableAttr(attr.getName(), attr.getValue()); + } + cachedReads.emplace_back(readAttrs, scalarRead.getResult()); + } + scalarValueByRead[readOp] = scalarRead.getResult(); + } + + for (ReadArrayOp readOp : llvm::make_early_inc_range(arrayReads)) { + replaceAllUsesIgnoringType(readOp.getResult(), scalarValueByRead.lookup(readOp)); + rewriter.eraseOp(readOp); + } + if (memberReadOp.getResult().use_empty()) { + rewriter.eraseOp(memberReadOp); + } + } + return success(); +} + +/// Reconstruct the scalar member type map established by the collected candidates. +static LogicalResult collectSplitTypesByMember( + ArrayRef arraysToScalarize, SymbolTableCollection &tables, + DenseMap> &splitIndicesByMember, + DenseMap> &splitTypesByMember, + const ConversionTracker &tracker +) { + for (const ScalarizedArrayInfo &info : arraysToScalarize) { + for (MemberWriteOp memberWriteOp : info.memberWrites) { + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + MemberDefOp member = memberDef->get(); + appendMissingIndices(splitIndicesByMember[member], info.indices); + DenseMap &splitTypes = splitTypesByMember[member]; + for (ArrayAttr idx : info.indices) { + auto res = mergeSplitCandidateType( + member, idx, info.typeByIndex.lookup(idx), memberWriteOp.getOperation(), splitTypes, + nullptr, tracker + ); + if (failed(res)) { + return failure(); + } + } + } + } + return success(); +} + +/// Verify one read of an array-typed member before the defining member is split. +static LogicalResult verifySplitMemberReadRewritable( + MemberReadOp memberReadOp, MemberDefOp member, const DenseMap &splitTypes, + SymbolTableCollection &tables, const ConversionTracker &tracker +) { + for (Operation *user : memberReadOp.getResult().getUsers()) { + auto readOp = llvm::dyn_cast(user); + if (!readOp || readOp.getArrRef() != memberReadOp.getResult()) { + InFlightDiagnostic diag = memberReadOp.emitError( + "cannot split heterogeneous array member because a whole-array read has an " + "unsupported use" + ); + diag.attachNote(member.getLoc()) << "split member defined here"; + diag.attachNote(user->getLoc()) << "unsupported use is here"; + return diag; + } + + ArrayAttr idx = getIndexAsAttr(readOp); + if (!idx) { + InFlightDiagnostic diag = readOp.emitError( + "cannot split heterogeneous array member because a read of it uses a dynamic array index" + ); + diag.attachNote(member.getLoc()) << "split member defined here"; + return diag; + } + + Type replacementType = splitTypes.lookup(idx); + if (!replacementType) { + InFlightDiagnostic diag = readOp.emitError( + "cannot split heterogeneous array member because a read of it uses an array index " + "that is not written by any scalarization candidate" + ); + diag.attachNote(member.getLoc()) << "split member defined here"; + return diag; + } + + if (!canReplaceReadResultWithType( + readOp, replacementType, tracker, "verifySplitMemberReadsRewritable" + )) { + InFlightDiagnostic diag = readOp.emitError( + "cannot split heterogeneous array member because a read result type is incompatible " + "with the split scalar member type" + ); + diag.attachNote(member.getLoc()) << "split member defined here"; + return diag; + } + + if (failed(canReplaceReadUsersWithType(readOp, replacementType, tables))) { + return failure(); + } + } + return success(); +} + +/// Verify reads of array-typed members before their defining members are split. +static LogicalResult verifySplitMemberReadsRewritable( + ModuleOp modOp, const DenseMap> &splitTypesByMember, + SymbolTableCollection &tables, const ConversionTracker &tracker +) { + for (MemberReadOp memberReadOp : walkCollect(modOp)) { + auto memberDef = memberReadOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + auto splitIt = splitTypesByMember.find(memberDef->get()); + if (splitIt == splitTypesByMember.end()) { + continue; + } + auto verifyRes = verifySplitMemberReadRewritable( + memberReadOp, memberDef->get(), splitIt->second, tables, tracker + ); + if (failed(verifyRes)) { + return failure(); + } + } + return success(); +} + +/// Return the local array allocation stored by `memberWriteOp`, if it has a static shape. +static FailureOr getStaticLocalArrayCreate(MemberWriteOp memberWriteOp) { + auto createOp = memberWriteOp.getVal().getDefiningOp(); + if (!createOp) { + return failure(); + } + ArrayType arrTy = createOp.getType(); + if (!arrTy.hasStaticShape() || llvm::isa(arrTy.getElementType())) { + return failure(); + } + return createOp; +} + +/// Verify that `createOp` can be expanded while splitting one or more members. +/// +/// Static writes update the local value map, static reads consume it, and missing indices are +/// materialized as shared `llzk.nondet` values. Other users would observe or escape the array in a +/// way this local scalarization cannot preserve. +static LogicalResult verifyExpandableLocalArrayUsers( + CreateArrayOp createOp, const DenseSet &splitMemberSet, + SymbolTableCollection &tables, const DominanceInfo &domInfo +) { + Value arrayValue = createOp.getResult(); + ArrayType arrTy = createOp.getType(); + std::optional> maybeIndices = arrTy.getSubelementIndices(); + if (!maybeIndices) { + return failure(); + } + bool allowCrossBlockMemberWrites = isInitializerOnlyLocalArray(createOp, *maybeIndices); + for (Operation *user : createOp.getResult().getUsers()) { + if (auto writeOp = llvm::dyn_cast(user)) { + if (writeOp.getArrRef() != arrayValue || writeOp->getBlock() != createOp->getBlock() || + !getIndexAsAttr(writeOp)) { + return failure(); + } + continue; + } + if (auto readOp = llvm::dyn_cast(user)) { + if (readOp.getArrRef() != arrayValue || readOp->getBlock() != createOp->getBlock() || + !getIndexAsAttr(readOp)) { + return failure(); + } + continue; + } + if (auto userWriteOp = llvm::dyn_cast(user)) { + auto memberDef = userWriteOp.getMemberDefOp(tables); + if (userWriteOp.getVal() == arrayValue && succeeded(memberDef) && + splitMemberSet.contains(memberDef->get())) { + if (userWriteOp->getBlock() == createOp->getBlock()) { + continue; + } + if (allowCrossBlockMemberWrites && + initializerValuesDominateUse(createOp, userWriteOp, domInfo)) { + continue; + } + } + } + return failure(); + } + return success(); +} + +/// Verify that every split-member write fed by `createOp` has the same static index set as the +/// local array allocation. +/// +/// `rewriteExpandableLocalArray()` emits one scalar write per target split index. If the source +/// allocation has extra indices, those values would be dropped; if it has missing indices, fresh +/// nondeterministic values would be materialized. Either changes the whole-array write semantics. +static LogicalResult verifyExpandableLocalArraySplitIndices( + CreateArrayOp createOp, + const DenseMap> &splitIndicesByMember, + const DenseMap *splitIndexOpsByMember, SymbolTableCollection &tables +) { + ArrayType arrTy = createOp.getType(); + std::optional> maybeIndices = arrTy.getSubelementIndices(); + if (!maybeIndices) { + return failure(); + } + ArrayRef indices = *maybeIndices; + Value arrayValue = createOp.getResult(); + + for (Operation *user : arrayValue.getUsers()) { + auto memberWriteOp = llvm::dyn_cast(user); + if (!memberWriteOp || memberWriteOp.getVal() != arrayValue) { + continue; + } + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + MemberDefOp member = memberDef->get(); + auto splitIt = splitIndicesByMember.find(member); + if (splitIt == splitIndicesByMember.end()) { + continue; + } + ArrayRef splitIndices = splitIt->second; + if (haveSameIndexSet(indices, splitIndices)) { + continue; + } + + InFlightDiagnostic diag = member.emitError( + "cannot split heterogeneous array member because expandable whole-array write uses " + "different index set" + ); + if (splitIndexOpsByMember) { + diag.attachNote(splitIndexOpsByMember->lookup(member)->getLoc()) + << "candidate establishes " << splitIndices.size() << " split member indices"; + } + Diagnostic ¬e = diag.attachNote(memberWriteOp.getLoc()) + << "expandable whole-array write has " << indices.size() << " array indices"; + if (ArrayAttr extraIndex = findIndexMissingFrom(indices, splitIndices)) { + note << ", including extra index " << extraIndex; + } else if (ArrayAttr missingIndex = findIndexMissingFrom(splitIndices, indices)) { + note << ", missing index " << missingIndex; + } + return diag; + } + return success(); +} +/// Reject expandable arrays that would store one SSA value into several split scalar members that +/// cannot share one materialized value. +/// +/// Later type propagation may refine the value operand of each scalar member write. If two writes +/// keep the same SSA value, refining one write also refines the other because MLIR values carry one +/// global type. Missing values are materialized as fresh `llzk.nondet` values during rewriting and +/// then cached by index, so the check tracks those pending materializations until an `array.write` +/// overwrites the index. Static reads of unwritten indices participate in the same pending +/// materialization state because the rewrite caches their result-typed `llzk.nondet` value for +/// later scalar member writes at that index. +static LogicalResult verifyNoSharedValuesForIncompatibleSplitTypes( + CreateArrayOp createOp, + const DenseMap> &splitTypesByMember, + SymbolTableCollection &tables, const ConversionTracker &tracker, + StringRef arrayDescription = "an expandable array" +) { + if (splitTypesByMember.empty()) { + return success(); + } + + bool targetsSplitMember = false; + for (Operation *user : createOp.getResult().getUsers()) { + auto memberWriteOp = llvm::dyn_cast(user); + if (!memberWriteOp || memberWriteOp.getVal() != createOp.getResult()) { + continue; + } + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + if (splitTypesByMember.contains(memberDef->get())) { + targetsSplitMember = true; + break; + } + } + if (!targetsSplitMember) { + return success(); + } + + ArrayType arrTy = createOp.getType(); + std::optional> maybeIndices = arrTy.getSubelementIndices(); + if (!maybeIndices) { + return failure(); + } + ArrayRef indices = *maybeIndices; + + DenseMap valueByIndex; + if (failed(seedValuesFromArrayElements(createOp, indices, valueByIndex))) { + return failure(); + } + + FailureOr> maybeUsers = getUsersInBlockOrder(createOp.getResult()); + SmallVector users; + if (succeeded(maybeUsers)) { + users = std::move(*maybeUsers); + } else { + if (hasDirectArrayWrites(createOp)) { + return failure(); + } + users.assign(createOp.getResult().user_begin(), createOp.getResult().user_end()); + } + + /// First consumer that forced an array index to be materialized at a concrete type. + /// + /// Rewriting creates and caches one `llzk.nondet` per unwritten index. It also caches the + /// refined type of a read from a stored `llzk.nondet` value. Recording the first requested type + /// lets verification reject a later read or split-member write that would need the same cached + /// value at an incompatible type. A stored `llzk.nondet` can be materialized at a common + /// refinement, as can a `poly.unifiable_cast` whose input can supply that refinement. Other + /// non-nondeterministic stored values cannot serve distinct split-member types. + struct PendingMaterializationUse { + /// Concrete type requested for an index before a write defines a new value there. + Type type; + /// Operation location used to explain the first pending materialization request. + Location loc; + /// Short diagnostic phrase naming the consumer kind that requested `type`. + StringRef description; + /// Diagnostic phrase naming whether this was an unwritten or already-stored index. + StringRef indexDescription; + }; + + DenseMap> firstUseByValue; + DenseMap firstMaterializedUseByIndex; + DenseMap> readsByIndex; + auto valueCompatibleWithTargetType = [&tracker](Value scalarValue, Type targetType) { + return canUseScalarizedValueAsType( + scalarValue.getType(), targetType, tracker, "verifyNoSharedValuesForIncompatibleSplitTypes" + ); + }; + auto verifyExternalMemberWriteUses = [&](Value scalarValue, ArrayAttr idx, Type targetType, + Location targetLoc) -> LogicalResult { + for (OpOperand &use : scalarValue.getUses()) { + auto externalWrite = llvm::dyn_cast(use.getOwner()); + if (!externalWrite || externalWrite.getVal() != use.get()) { + continue; + } + auto memberDef = externalWrite.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + Type externalType = memberDef->get().getType(); + if (succeeded(getCommonRefinedType(externalType, targetType, tracker))) { + continue; + } + + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription << " reuses one SSA value for incompatible scalar member types"; + diag.attachNote(externalWrite.getLoc()) << "value is used for member type " << externalType; + diag.attachNote(targetLoc) << "same value is also used for array index " << idx + << " scalar member type " << targetType; + return diag; + } + return success(); + }; + auto noteMaterializedUse = [&](ArrayAttr idx, Type targetType, Location loc, + StringRef description, + StringRef indexDescription) -> LogicalResult { + auto existing = firstMaterializedUseByIndex.find(idx); + if (existing == firstMaterializedUseByIndex.end()) { + firstMaterializedUseByIndex.try_emplace( + idx, PendingMaterializationUse {targetType, loc, description, indexDescription} + ); + return success(); + } + Type existingType = existing->second.type; + FailureOr commonType = getCommonRefinedType(existingType, targetType, tracker); + if (failed(commonType)) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription << " reuses one SSA value for incompatible scalar member types"; + diag.attachNote(existing->second.loc) + << existing->second.indexDescription << ' ' << idx << " is materialized for " + << existing->second.description << ' ' << existingType; + diag.attachNote(loc) << "same index is also materialized for " << description << ' ' + << targetType; + return diag; + } + existing->second.type = *commonType; + return success(); + }; + // The rewrite replaces every read in one array-value lifetime with the same materialized + // scalar value. Check the final type against both the read result and its typed consumers before + // any candidate or expandable array rewrite can erase an earlier read. + auto verifyMaterializedReadUsers = [&](ArrayAttr idx) -> LogicalResult { + auto materialization = firstMaterializedUseByIndex.find(idx); + if (materialization == firstMaterializedUseByIndex.end()) { + return success(); + } + for (ReadArrayOp readOp : readsByIndex.lookup(idx)) { + if (!canReplaceReadResultWithType( + readOp, materialization->second.type, tracker, + "verifyNoSharedValuesForIncompatibleSplitTypes" + ) || + failed(canReplaceReadUsersWithType(readOp, materialization->second.type, tables))) { + return failure(); + } + } + readsByIndex.erase(idx); + firstMaterializedUseByIndex.erase(idx); + return success(); + }; + for (Operation *user : users) { + if (auto writeOp = llvm::dyn_cast(user)) { + ArrayAttr idx = getIndexAsAttr(writeOp); + if (failed(verifyMaterializedReadUsers(idx))) { + return failure(); + } + valueByIndex[idx] = writeOp.getRvalue(); + continue; + } + if (auto readOp = llvm::dyn_cast(user)) { + ArrayAttr idx = getIndexAsAttr(readOp); + Type readType = readOp.getResult().getType(); + if (Value scalarValue = valueByIndex.lookup(idx)) { + if (failed(getCommonRefinedType(scalarValue.getType(), readType, tracker))) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription << " stores a scalar value with an incompatible read result type"; + diag.attachNote(scalarValue.getLoc()) + << "array index " << idx << " stores value type " << scalarValue.getType(); + diag.attachNote(readOp.getLoc()) << "same index is read as " << readType; + return diag; + } + auto res = + noteMaterializedUse(idx, readType, readOp.getLoc(), "read result type", "array index"); + if (failed(res)) { + return failure(); + } + } else { + auto res = noteMaterializedUse( + idx, readType, readOp.getLoc(), "read result type", "unwritten index" + ); + if (failed(res)) { + return failure(); + } + } + readsByIndex[idx].push_back(readOp); + continue; + } + + auto memberWriteOp = llvm::dyn_cast(user); + if (!memberWriteOp) { + continue; + } + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + auto splitIt = splitTypesByMember.find(memberDef->get()); + if (splitIt == splitTypesByMember.end()) { + continue; + } + + for (const auto &[idx, targetType] : splitIt->second) { + Value scalarValue = valueByIndex.lookup(idx); + if (!scalarValue) { + auto res = noteMaterializedUse( + idx, targetType, memberWriteOp.getLoc(), "scalar member type", "unwritten index" + ); + if (failed(res)) { + return failure(); + } + continue; + } + auto existing = firstUseByValue.find(scalarValue); + if (!valueCompatibleWithTargetType(scalarValue, targetType)) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription << " stores a scalar value with an incompatible split-member type"; + diag.attachNote(scalarValue.getLoc()) + << "array index " << idx << " stores value type " << scalarValue.getType(); + diag.attachNote(memberWriteOp.getLoc()) + << "same index is written with scalar member type " << targetType; + return diag; + } + auto verifyWritesRes = + verifyExternalMemberWriteUses(scalarValue, idx, targetType, memberWriteOp.getLoc()); + if (failed(verifyWritesRes)) { + return failure(); + } + auto noteMaterializedUseRes = noteMaterializedUse( + idx, targetType, memberWriteOp.getLoc(), "scalar member type", "array index" + ); + if (failed(noteMaterializedUseRes)) { + return failure(); + } + if (existing == firstUseByValue.end()) { + bool canMaterializeTarget; + if (UnifiableCastOp castOp = scalarValue.getDefiningOp()) { + canMaterializeTarget = canUseScalarizedValueAsType( + castOp.getInput().getType(), targetType, tracker, + "verifyNoSharedValuesForIncompatibleSplitTypes" + ); + } else { + // Array reads are replaced before member writes, so their replacement can carry the + // refined type. Other ordinary SSA values would be emitted unchanged. + canMaterializeTarget = + scalarValue.getDefiningOp() || scalarValue.getDefiningOp(); + } + if (targetType != scalarValue.getType() && !canMaterializeTarget) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription + << " uses a non-nondeterministic SSA value for a refined scalar member type"; + diag.attachNote(scalarValue.getLoc()) + << "array index " << idx << " stores value type " << scalarValue.getType(); + diag.attachNote(memberWriteOp.getLoc()) + << "same index is written with scalar member type " << targetType; + return diag; + } + firstUseByValue.try_emplace( + scalarValue, std::make_pair(targetType, memberWriteOp.getLoc()) + ); + continue; + } + Type existingType = existing->second.first; + FailureOr commonType = getCommonRefinedType(existingType, targetType, tracker); + if (failed(commonType)) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription << " reuses one SSA value for incompatible scalar member types"; + diag.attachNote(existing->second.second) + << "value is used for scalar member type " << existingType; + diag.attachNote(memberWriteOp.getLoc()) + << "same value is also used for scalar member type " << targetType; + return diag; + } + bool canMaterializeCommonType; + if (UnifiableCastOp castOp = scalarValue.getDefiningOp()) { + canMaterializeCommonType = canUseScalarizedValueAsType( + castOp.getInput().getType(), *commonType, tracker, + "verifyNoSharedValuesForIncompatibleSplitTypes" + ); + } else { + canMaterializeCommonType = scalarValue.getDefiningOp(); + } + if (*commonType != scalarValue.getType() && !canMaterializeCommonType) { + InFlightDiagnostic diag = + createOp.emitError("cannot split heterogeneous array member because ") + << arrayDescription + << " reuses a non-nondeterministic SSA value for distinct scalar member types"; + diag.attachNote(existing->second.second) + << "value is used for scalar member type " << existingType; + diag.attachNote(memberWriteOp.getLoc()) + << "same value is also used for scalar member type " << targetType; + return diag; + } + existing->second.first = *commonType; + } + } + for (ArrayAttr idx : indices) { + if (failed(verifyMaterializedReadUsers(idx))) { + return failure(); + } + } + return success(); +} + +/// Rewrite one expandable local array allocation in block order. +static LogicalResult rewriteExpandableLocalArray( + CreateArrayOp createOp, const DenseMap &splitMembers, + SymbolTableCollection &tables, PatternRewriter &rewriter, const ConversionTracker &tracker +) { + ArrayType arrTy = createOp.getType(); + std::optional> maybeIndices = arrTy.getSubelementIndices(); + if (!maybeIndices) { + return failure(); + } + ArrayRef indices = *maybeIndices; + + FailureOr> maybeUsers = getUsersInBlockOrder(createOp.getResult()); + Operation *materializationPoint = nullptr; + SmallVector users; + if (succeeded(maybeUsers)) { + users = std::move(*maybeUsers); + } else { + if (!isInitializerOnlyLocalArray(createOp, indices)) { + return failure(); + } + users.assign(createOp.getResult().user_begin(), createOp.getResult().user_end()); + materializationPoint = createOp.getOperation(); + } + + DenseMap valueByIndex; + DenseMap writeDiscardableAttrsByIndex; + if (failed(seedValuesFromArrayElements(createOp, indices, valueByIndex))) { + return failure(); + } + + // Plan one final type for every materialization before rewriting any reads. Reads of an + // unwritten index share a cached nondet value until an array write replaces it. In particular, + // a later complementary generic read can refine that cached value, so validating only the + // earlier read's initial replacement type can leave its consumers invalid after the refinement. + DenseMap materializationTypeByRead; + DenseMap materializationTypeByIndex; + DenseMap> readsByIndex; + auto addMaterializationType = [&](ArrayAttr idx, Type type) -> LogicalResult { + auto existing = materializationTypeByIndex.find(idx); + if (existing == materializationTypeByIndex.end()) { + materializationTypeByIndex.try_emplace(idx, type); + return success(); + } + FailureOr commonType = getCommonRefinedType(existing->second, type, tracker); + if (failed(commonType)) { + return failure(); + } + existing->second = *commonType; + return success(); + }; + auto finishMaterialization = [&](ArrayAttr idx) { + Type finalType = materializationTypeByIndex.lookup(idx); + if (finalType) { + for (Operation *read : readsByIndex.lookup(idx)) { + materializationTypeByRead[read] = finalType; + } + } + materializationTypeByIndex.erase(idx); + readsByIndex.erase(idx); + }; + for (Operation *user : users) { + if (auto writeOp = llvm::dyn_cast(user)) { + finishMaterialization(getIndexAsAttr(writeOp)); + continue; + } + if (auto readOp = llvm::dyn_cast(user)) { + ArrayAttr idx = getIndexAsAttr(readOp); + if (failed(addMaterializationType(idx, readOp.getResult().getType()))) { + return failure(); + } + readsByIndex[idx].push_back(readOp); + continue; + } + auto memberWriteOp = llvm::dyn_cast(user); + if (!memberWriteOp) { + continue; + } + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + auto splitIt = splitMembers.find(memberDef->get()); + if (splitIt == splitMembers.end()) { + return failure(); + } + for (ArrayAttr idx : splitIt->second.indices) { + MemberInfo memberInfo = splitIt->second.memberByIndex.lookup(idx); + if (!memberInfo.first || failed(addMaterializationType(idx, memberInfo.second))) { + return failure(); + } + } + } + for (ArrayAttr idx : indices) { + finishMaterialization(idx); + } + + // Validate every consumer against the final shared type before replacing any read. This keeps a + // later refinement from changing the type of a nondet value that an earlier replacement already + // made visible to an incompatible consumer. + for (const auto &[read, finalType] : materializationTypeByRead) { + if (failed(canReplaceReadUsersWithType(llvm::cast(read), finalType, tables))) { + return failure(); + } + } + DenseSet generatedMaterializations; + SmallVector writesToErase; + for (Operation *user : users) { + if (auto writeOp = llvm::dyn_cast(user)) { + ArrayAttr idx = getIndexAsAttr(writeOp); + valueByIndex[idx] = writeOp.getRvalue(); + writeDiscardableAttrsByIndex[idx] = writeOp->getDiscardableAttrDictionary(); + writesToErase.push_back(writeOp); + continue; + } + + if (auto readOp = llvm::dyn_cast(user)) { + ArrayAttr idx = getIndexAsAttr(readOp); + Type requestedType = readOp.getResult().getType(); + if (Type materializationType = materializationTypeByRead.lookup(readOp.getOperation())) { + requestedType = materializationType; + } + rewriter.setInsertionPoint(readOp); + FailureOr scalarValue = getOrCreateScalarizedLocalArrayValue( + valueByIndex, generatedMaterializations, idx, requestedType, readOp.getLoc(), rewriter, + tracker, materializationPoint + ); + if (failed(scalarValue)) { + return failure(); + } + replaceAllUsesIgnoringType(readOp.getResult(), *scalarValue); + rewriter.eraseOp(readOp); + continue; + } + + auto memberWriteOp = llvm::dyn_cast(user); + if (!memberWriteOp) { + return failure(); + } + + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (failed(memberDef)) { + return failure(); + } + auto splitIt = splitMembers.find(memberDef->get()); + if (splitIt == splitMembers.end()) { + return failure(); + } + const SplitMemberInfo &splitInfo = splitIt->second; + + auto getScalarValue = [&](ArrayAttr idx, Type type) { + return getOrCreateScalarizedLocalArrayValue( + valueByIndex, generatedMaterializations, idx, type, memberWriteOp.getLoc(), rewriter, + tracker, materializationPoint + ); + }; + auto emitRes = emitScalarMemberWrites( + memberWriteOp, splitInfo, writeDiscardableAttrsByIndex, rewriter, getScalarValue + ); + if (failed(emitRes)) { + return failure(); + } + rewriter.eraseOp(memberWriteOp); + } + + for (WriteArrayOp writeOp : llvm::make_early_inc_range(writesToErase)) { + rewriter.eraseOp(writeOp); + } + if (createOp.getResult().use_empty()) { + rewriter.eraseOp(createOp); + } + return success(); +} + +/// Rewrite remaining expandable whole-array writes to split members. +static LogicalResult rewriteExpandableMemberWrites( + DenseMap &splitMembers, SymbolTableCollection &tables, + PatternRewriter &rewriter, const ConversionTracker &tracker +) { + DenseSet splitMemberSet; + DenseMap> splitIndicesByMember; + for (const auto &entry : splitMembers) { + splitMemberSet.insert(entry.first); + splitIndicesByMember[entry.first] = entry.second.indices; + } + + for (const auto &entry : splitMembers) { + MemberDefOp member = entry.first; + StructDefOp parentStruct = getParentOfType(member); + assert(parentStruct && "MemberDefOp parent is always StructDefOp"); + + auto uses = llzk::getSymbolUses(member, parentStruct); + if (!uses) { + return failure(); + } + + SmallVector createsToExpand; + DenseSet seenCreates; + for (SymbolTable::SymbolUse symUse : uses.value()) { + auto memberWriteOp = llvm::dyn_cast(symUse.getUser()); + if (!memberWriteOp) { + continue; + } + auto memberDef = memberWriteOp.getMemberDefOp(tables); + FailureOr maybeCreateOp = getStaticLocalArrayCreate(memberWriteOp); + if (succeeded(memberDef) && memberDef->get() == member && succeeded(maybeCreateOp) && + !seenCreates.contains(*maybeCreateOp)) { + createsToExpand.push_back(*maybeCreateOp); + seenCreates.insert(*maybeCreateOp); + } + } + + for (CreateArrayOp createOp : createsToExpand) { + DominanceInfo domInfo; + if (failed(verifyExpandableLocalArrayUsers(createOp, splitMemberSet, tables, domInfo))) { + return failure(); + } + auto res = verifyExpandableLocalArraySplitIndices( + createOp, splitIndicesByMember, /*splitIndexOpsByMember=*/nullptr, tables + ); + if (failed(res)) { + return failure(); + } + if (failed(rewriteExpandableLocalArray(createOp, splitMembers, tables, rewriter, tracker))) { + return failure(); + } + } + } + return success(); +} + +/// Verify that all writes to any split member are either candidate or otherwise expandable. +/// +/// Splitting a member is global: after replacement, every read of the original array-typed member +/// is redirected to the split scalar members. That is only sound when every write to that member +/// can also be expanded. +static LogicalResult verifySplitMemberWritesExpandable( + ArrayRef arraysToScalarize, SymbolTableCollection &tables, + const ConversionTracker &tracker +) { + DenseMap> candidateWritesByMember; + DenseMap> splitIndicesByMember; + DenseMap splitIndexOpsByMember; + DenseMap> splitTypesByMember; + DenseMap> splitTypeOpsByMember; + for (const ScalarizedArrayInfo &info : arraysToScalarize) { + for (MemberWriteOp memberWriteOp : info.memberWrites) { + auto memberDef = memberWriteOp.getMemberDefOp(tables); + if (succeeded(memberDef)) { + MemberDefOp member = memberDef->get(); + candidateWritesByMember[member].insert(memberWriteOp.getOperation()); + auto [indicesIt, inserted] = splitIndicesByMember.try_emplace(member, info.indices); + if (inserted) { + splitIndexOpsByMember[member] = memberWriteOp.getOperation(); + } else if (!haveSameIndexSet(indicesIt->second, info.indices)) { + InFlightDiagnostic diag = member.emitError( + "cannot split heterogeneous array member because candidate whole-array writes use " + "different index set" + ); + diag.attachNote(splitIndexOpsByMember.lookup(member)->getLoc()) + << "candidate establishes " << indicesIt->second.size() << " split member indices"; + Diagnostic ¬e = diag.attachNote(memberWriteOp.getLoc()) + << "candidate whole-array write has " << info.indices.size() + << " array indices"; + if (ArrayAttr extraIndex = findIndexMissingFrom(info.indices, indicesIt->second)) { + note << ", including extra index " << extraIndex; + } else if (ArrayAttr missingIndex = + findIndexMissingFrom(indicesIt->second, info.indices)) { + note << ", missing index " << missingIndex; + } + return diag; + } + DenseMap &splitTypes = splitTypesByMember[member]; + DenseMap &splitTypeOps = splitTypeOpsByMember[member]; + for (ArrayAttr idx : info.indices) { + Type candidateType = info.typeByIndex.lookup(idx); + auto res = mergeSplitCandidateType( + member, idx, candidateType, memberWriteOp.getOperation(), splitTypes, &splitTypeOps, + tracker + ); + if (failed(res)) { + return failure(); + } + } + } + } + } + DenseSet splitMemberSet; + for (const auto &entry : candidateWritesByMember) { + splitMemberSet.insert(entry.first); + } + DominanceInfo domInfo; + + for (const ScalarizedArrayInfo &info : arraysToScalarize) { + auto verifyNoIncompatibleShares = verifyNoSharedValuesForIncompatibleSplitTypes( + info.createOp, splitTypesByMember, tables, tracker, "a scalarization candidate" + ); + if (failed(verifyNoIncompatibleShares)) { + return failure(); + } + } + + for (const auto &entry : candidateWritesByMember) { + MemberDefOp member = entry.first; + const DenseSet &candidateWrites = entry.second; + StructDefOp parentStruct = getParentOfType(member); + assert(parentStruct && "MemberDefOp parent is always StructDefOp"); + + auto uses = llzk::getSymbolUses(member, parentStruct); + if (!uses) { + return member.emitError( + "cannot split heterogeneous array member because its symbol uses could not be inspected" + ); + } + + for (SymbolTable::SymbolUse symUse : uses.value()) { + auto writeOp = llvm::dyn_cast(symUse.getUser()); + if (!writeOp || candidateWrites.contains(writeOp.getOperation())) { + continue; + } + FailureOr maybeCreateOp = getStaticLocalArrayCreate(writeOp); + if (failed(maybeCreateOp) || failed(verifyExpandableLocalArrayUsers( + *maybeCreateOp, splitMemberSet, tables, domInfo + ))) { + InFlightDiagnostic diag = member.emitError( + "cannot split heterogeneous array member because not every write to it can be " + "scalarized" + ); + diag.attachNote(writeOp.getLoc()) << "whole-array write is not backed by a scalarization " + "candidate"; + return diag; + } + auto verifyExpandable = verifyExpandableLocalArraySplitIndices( + *maybeCreateOp, splitIndicesByMember, &splitIndexOpsByMember, tables + ); + if (failed(verifyExpandable)) { + return failure(); + } + auto verifyNoIncompatibleShares = verifyNoSharedValuesForIncompatibleSplitTypes( + *maybeCreateOp, splitTypesByMember, tables, tracker + ); + if (failed(verifyNoIncompatibleShares)) { + return failure(); + } + } + } + return success(); +} + +/// Erase original array-typed members after all symbol uses have been redirected. +static void eraseUnusedOriginalMembers( + DenseMap &splitMembers, PatternRewriter &rewriter +) { + for (const auto &[member, _] : splitMembers) { + StructDefOp parentStruct = getParentOfType(member); + assert(parentStruct && "MemberDefOp parent is always StructDefOp"); + auto uses = llzk::getSymbolUses(member, parentStruct); + if (uses && uses->empty()) { + rewriter.eraseOp(member); + } + } +} + +/// Scalarize all safe pseudo-homogeneous arrays exposed in the current flattening iteration. +/// +/// Running this before general type propagation prevents the propagation step from choosing one +/// concrete element type for an array that semantically contains a different concrete type at each +/// static index. +LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { + SmallVector arraysToScalarize; + modOp.walk([&arraysToScalarize, &tracker](CreateArrayOp op) { + FailureOr info = getScalarizedArrayInfo(op, tracker); + if (succeeded(info)) { + arraysToScalarize.push_back(*info); + } + }); + if (arraysToScalarize.empty()) { + return success(); + } + + PatternRewriter rewriter(modOp.getContext()); + SymbolTableCollection tables; + DenseMap> splitIndicesByMember; + DenseMap> splitTypesByMember; + DenseMap candidateInfoByCreateOp; + auto refreshRes = + refreshValuesFromDependentCandidates(arraysToScalarize, candidateInfoByCreateOp, tracker); + if (failed(refreshRes)) { + return failure(); + } + auto verifyRes = + verifyRefreshedCandidateReads(arraysToScalarize, tables, tracker, candidateInfoByCreateOp); + if (failed(verifyRes)) { + return failure(); + } + if (failed(verifySplitMemberWritesExpandable(arraysToScalarize, tables, tracker))) { + return failure(); + } + auto collectRes = collectSplitTypesByMember( + arraysToScalarize, tables, splitIndicesByMember, splitTypesByMember, tracker + ); + if (failed(collectRes) || + failed(verifySplitMemberReadsRewritable(modOp, splitTypesByMember, tables, tracker))) { + return failure(); + } + for (const ScalarizedArrayInfo &info : arraysToScalarize) { + DenseMap specializedTypeByNondet; + auto collectSpecializedRes = collectSharedNondetSpecializationTypes( + info, tracker, splitTypesByMember, tables, specializedTypeByNondet + ); + if (failed(collectSpecializedRes)) { + return failure(); + } + auto verifySpecializedRes = + verifySharedNondetSpecializationExternalUses(info, specializedTypeByNondet, tables); + if (failed(verifySpecializedRes)) { + return failure(); + } + } + + DenseMap splitMembers; + for (ScalarizedArrayInfo &info : arraysToScalarize) { + auto rewriteRes = rewriteLocalArray( + info, splitMembers, splitTypesByMember, splitIndicesByMember, tables, rewriter, tracker, + &candidateInfoByCreateOp + ); + if (failed(rewriteRes)) { + return failure(); + } + } + if (failed(rewriteExpandableMemberWrites(splitMembers, tables, rewriter, tracker))) { + return failure(); + } + if (failed(rewriteSplitMemberReads(modOp, splitMembers, tables, rewriter, tracker))) { + return failure(); + } + eraseUnusedOriginalMembers(splitMembers, rewriter); + tracker.updateModifiedFlag(true); + return success(); +} + +} // namespace Step5_ScalarizeHeterogeneousArrays + +namespace Step6_PropagateTypes { + +/// Update the array element type from compatible initializers and writes. +class UpdateNewArrayElemFromWrite final : public OpRewritePattern { + ConversionTracker &tracker_; + +public: + /// Construct the create-array element-type propagation pattern. + UpdateNewArrayElemFromWrite(MLIRContext *ctx, ConversionTracker &tracker) + : OpRewritePattern(ctx, 3), tracker_(tracker) {} + + /// Update an `array.new` result element type from compatible initializers and writes. + LogicalResult matchAndRewrite(CreateArrayOp op, PatternRewriter &rewriter) const override { + Value createResult = op.getResult(); + ArrayType createResultType = dyn_cast(createResult.getType()); + assert(createResultType && "CreateArrayOp must produce ArrayType"); + Type oldResultElemType = createResultType.getElementType(); + + Type newResultElemType = nullptr; + // An initializer constrains array.new just as strongly as a later array.write: every element + // must have the result array's element type. A member-read refinement can change an + // initializer without creating a WriteArrayOp, so consider the explicit elements first. + if (!op.getElements().empty()) { + Type initializerType = op.getElements().front().getType(); + if (!llvm::all_of(op.getElements(), [initializerType](Value element) { + return element.getType() == initializerType; + })) { + return failure(); + } + if (initializerType != oldResultElemType) { + newResultElemType = initializerType; + } + } + + // Look for WriteArrayOp where the array reference is the result of the CreateArrayOp and the + // element type is different. + for (Operation *user : createResult.getUsers()) { + if (WriteArrayOp writeOp = dyn_cast(user)) { + if (writeOp.getArrRef() != createResult) { + continue; + } + Type writeRValueType = writeOp.getRvalue().getType(); + if (writeRValueType == oldResultElemType) { + continue; + } + if (newResultElemType && newResultElemType != writeRValueType) { + LLVM_DEBUG( + llvm::dbgs() << "[UpdateNewArrayElemFromWrite] multiple possible element types for CreateArrayOp " << newResultElemType << " vs " << writeRValueType << '\n' ); @@ -2147,6 +5224,10 @@ class UpdateNewArrayElemFromWrite final : public OpRewritePattern return failure(); } ArrayType newType = createResultType.cloneWith(newResultElemType); + SymbolTableCollection tables; + if (!Step5_ScalarizeHeterogeneousArrays::canRefineCreateArrayUsersToType(op, newType, tables)) { + return failure(); + } rewriter.modifyOpInPlace(op, [&createResult, &newType]() { createResult.setType(newType); }); LLVM_DEBUG( llvm::dbgs() << "[UpdateNewArrayElemFromWrite] updated result type of " << op << '\n' @@ -2157,6 +5238,8 @@ class UpdateNewArrayElemFromWrite final : public OpRewritePattern namespace { +/// Update the array reference type on an array access op to match a scalar element type observed +/// through that access. LogicalResult updateArrayElemFromArrAccessOp( ArrayAccessOpInterface op, Type scalarElemTy, ConversionTracker &tracker, PatternRewriter &rewriter @@ -2183,9 +5266,11 @@ class UpdateArrayElemFromArrWrite final : public OpRewritePattern ConversionTracker &tracker_; public: + /// Construct the array-write based element-type propagation pattern. UpdateArrayElemFromArrWrite(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update the referenced array type from the write value type. LogicalResult matchAndRewrite(WriteArrayOp op, PatternRewriter &rewriter) const override { return updateArrayElemFromArrAccessOp(op, op.getRvalue().getType(), tracker_, rewriter); } @@ -2195,9 +5280,11 @@ class UpdateArrayElemFromArrRead final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the array-read based element-type propagation pattern. UpdateArrayElemFromArrRead(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update the referenced array type from the read result type. LogicalResult matchAndRewrite(ReadArrayOp op, PatternRewriter &rewriter) const override { return updateArrayElemFromArrAccessOp(op, op.getResult().getType(), tracker_, rewriter); } @@ -2208,9 +5295,11 @@ class UpdateMemberDefTypeFromWrite final : public OpRewritePattern ConversionTracker &tracker_; public: + /// Construct the member-definition propagation pattern. UpdateMemberDefTypeFromWrite(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update a member definition type from compatible writes to that member. LogicalResult matchAndRewrite(MemberDefOp op, PatternRewriter &rewriter) const override { // Find all uses of the member symbol name within its parent struct. StructDefOp parentRes = getParentOfType(op); @@ -2235,7 +5324,15 @@ class UpdateMemberDefTypeFromWrite final : public OpRewritePattern // A->B is a legal conversion (i.e., more concrete unification), then it is safe to use // type B with the assumption that the write with type A will be updated by another // pattern to also use type B. - if (!tracker_.isLegalConversion(writeToType, newType, "UpdateMemberDefTypeFromWrite")) { + auto commonType = Step5_ScalarizeHeterogeneousArrays::getCommonRefinedType( + newType, writeToType, tracker_ + ); + if (succeeded(commonType)) { + newType = *commonType; + newTypeLoc = writeOp.getLoc(); + } else if (!tracker_.isLegalConversion( + writeToType, newType, "UpdateMemberDefTypeFromWrite" + )) { if (tracker_.isLegalConversion( newType, writeToType, "UpdateMemberDefTypeFromWrite" )) { @@ -2265,7 +5362,34 @@ class UpdateMemberDefTypeFromWrite final : public OpRewritePattern return failure(); // nothing changed } if (!tracker_.isLegalConversion(op.getType(), newType, "UpdateMemberDefTypeFromWrite")) { - return failure(); + auto commonType = + Step5_ScalarizeHeterogeneousArrays::getCommonRefinedType(op.getType(), newType, tracker_); + if (failed(commonType) || *commonType != newType) { + return failure(); + } + } + + // Do not commit a write-derived refinement that an existing read or one of its typed consumers + // cannot adopt. In particular, two complementary writes can have a common refinement even + // when a generic read is independently consumed as an incompatible concrete type. + if (auto memberUsers = llzk::getSymbolUses(op, parentRes)) { + SymbolTableCollection tables; + for (SymbolTable::SymbolUse symUse : memberUsers.value()) { + if (MemberReadOp readOp = llvm::dyn_cast(symUse.getUser())) { + bool canUse = Step5_ScalarizeHeterogeneousArrays::canUseScalarizedValueAsType( + readOp.getVal().getType(), newType, tracker_, "UpdateMemberDefTypeFromWrite" + ); + if (!canUse) { + return failure(); + } + bool canRefine = Step5_ScalarizeHeterogeneousArrays::canRefineMemberReadUsersToType( + readOp, newType, tables, tracker_ + ); + if (!canRefine) { + return failure(); + } + } + } } rewriter.modifyOpInPlace(op, [&op, &newType]() { op.setType(newType); }); LLVM_DEBUG(llvm::dbgs() << "[UpdateMemberDefTypeFromWrite] updated type of " << op << '\n'); @@ -2275,6 +5399,7 @@ class UpdateMemberDefTypeFromWrite final : public OpRewritePattern namespace { +/// Move all regions out of `op` so it can be recreated with updated result types. SmallVector> moveRegions(Operation *op) { SmallVector> newRegions; for (Region ®ion : op->getRegions()) { @@ -2293,9 +5418,11 @@ class UpdateInferredResultTypes final : public OpTraitRewritePattern inferredResultTypes; InferTypeOpInterface retTypeFn = llvm::cast(op); @@ -2333,9 +5460,11 @@ class UpdateFuncTypeFromReturn final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the function-return based type propagation pattern. UpdateFuncTypeFromReturn(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update a function type from its terminator operand types. LogicalResult matchAndRewrite(FuncDefOp op, PatternRewriter &rewriter) const override { Region &body = op.getFunctionBody(); if (body.empty()) { @@ -2374,9 +5503,11 @@ class UpdateFreeFuncCallOpTypes final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the free-function call result propagation pattern. UpdateFreeFuncCallOpTypes(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Rewrite a call to a free function when the target function result types were refined. LogicalResult matchAndRewrite(CallOp op, PatternRewriter &rewriter) const override { if (calleeReferencesTemplateParam(op)) { return failure(); @@ -2411,6 +5542,7 @@ class UpdateFreeFuncCallOpTypes final : public OpRewritePattern { namespace { +/// Update a member read/write value type from the referenced member definition. LogicalResult updateMemberRefValFromMemberDef( MemberRefOpInterface op, ConversionTracker &tracker, PatternRewriter &rewriter ) { @@ -2422,9 +5554,18 @@ LogicalResult updateMemberRefValFromMemberDef( Type oldResultType = op.getVal().getType(); Type newResultType = def->get().getType(); if (oldResultType == newResultType || - !tracker.isLegalConversion(oldResultType, newResultType, "updateMemberRefValFromMemberDef")) { + !Step5_ScalarizeHeterogeneousArrays::canUseScalarizedValueAsType( + oldResultType, newResultType, tracker, "updateMemberRefValFromMemberDef" + )) { return failure(); } + if (MemberReadOp readOp = llvm::dyn_cast(op.getOperation())) { + if (!Step5_ScalarizeHeterogeneousArrays::canRefineMemberReadUsersToType( + readOp, newResultType, tables, tracker + )) { + return failure(); + } + } rewriter.modifyOpInPlace(op, [&op, &newResultType]() { op.getVal().setType(newResultType); }); LLVM_DEBUG( llvm::dbgs() << "[updateMemberRefValFromMemberDef] updated value type in " << op << '\n' @@ -2439,27 +5580,245 @@ class UpdateMemberReadValFromDef final : public OpRewritePattern { ConversionTracker &tracker_; public: + /// Construct the member-read value propagation pattern. UpdateMemberReadValFromDef(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update a member read result type from its referenced member definition. LogicalResult matchAndRewrite(MemberReadOp op, PatternRewriter &rewriter) const override { return updateMemberRefValFromMemberDef(op, tracker_, rewriter); } }; +namespace { + +/// Return whether `value` is used by any operation other than `user`. +static bool hasUsesOutside(Value value, Operation *user) { + return llvm::any_of(value.getUses(), [user](OpOperand &use) { return use.getOwner() != user; }); +} + +/// Return whether every use of `value` is a member-write value operand that accepts `type`. +/// +/// A generic cast result can feed several member writes whose member definitions are refined +/// independently. Retagging that shared result is valid only when every target member accepts the +/// cast input type. +static bool allMemberWritesAcceptType(Value value, Type type) { + if (value.use_empty()) { + return false; + } + + SymbolTableCollection tables; + return llvm::all_of(value.getUses(), [&tables, type](OpOperand &use) { + auto writeOp = llvm::dyn_cast(use.getOwner()); + if (!writeOp || writeOp.getVal() != use.get()) { + return false; + } + auto memberDef = writeOp.getMemberDefOp(tables); + return succeeded(memberDef) && typesUnify(type, memberDef->get().getType()); + }); +} + +/// Return a shared typed nondeterministic replacement for member writes that reuse `value`. +static FailureOr materializeValueForMemberWrite( + Location loc, Value value, Type type, PatternRewriter &rewriter, ConversionTracker &tracker, + DenseMap &typedNondetReplacements +) { + auto cachedReplacement = typedNondetReplacements.find(value); + if (cachedReplacement != typedNondetReplacements.end()) { + Value replacement = cachedReplacement->second; + FailureOr commonType = Step5_ScalarizeHeterogeneousArrays::getCommonRefinedType( + replacement.getType(), type, tracker + ); + if (failed(commonType)) { + return failure(); + } + if (*commonType != replacement.getType()) { + NonDetOp replacementNondetOp = replacement.getDefiningOp(); + if (!replacementNondetOp) { + return failure(); + } + rewriter.modifyOpInPlace(replacementNondetOp, [&replacement, &commonType]() { + replacement.setType(*commonType); + }); + } + return replacement; + } + if (value.getType() == type) { + return value; + } + + NonDetOp nondetOp = value.getDefiningOp(); + if (!nondetOp) { + return failure(); + } + + // A direct member write can become concrete after another write refines the member definition. + // Do not split a shared witness from a compatible external constraint while materializing the + // concrete value for that write. This mirrors the retargeting performed for scalarized array + // initializers: redirect the constraint to the instantiated struct definition and use the one + // replacement witness for both observations. + SmallVector externalConstraints; + SymbolTableCollection tables; + for (OpOperand &use : value.getUses()) { + auto memberWrite = llvm::dyn_cast(use.getOwner()); + if (memberWrite && memberWrite.getVal() == use.get()) { + continue; + } + auto call = llvm::dyn_cast(use.getOwner()); + if (!call || !call.calleeIsStructConstrain() || use.getOperandNumber() != 0 || + !llvm::isa(type)) { + return failure(); + } + + StructType structType = llvm::cast(type); + SymbolRefAttr callee = + appendLeaf(structType.getNameRef(), call.getCalleeAttr().getLeafReference()); + FailureOr> target = + lookupTopLevelSymbol(tables, callee, call); + if (failed(target)) { + return failure(); + } + SmallVector argTypes(call.getArgOperands().getTypes()); + argTypes.front() = type; + if (!typeListsUnify(argTypes, target->get().getArgumentTypes(), target->getNamespace())) { + return failure(); + } + externalConstraints.push_back(call); + } + + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(nondetOp); + auto clonedNondet = rewriter.create(loc, type); + clonedNondet->setDiscardableAttrs(nondetOp->getDiscardableAttrDictionary()); + Value replacement = clonedNondet.getResult(); + typedNondetReplacements.try_emplace(value, replacement); + for (CallOp call : externalConstraints) { + replaceStructConstraintCallWithSpecializedWitness(rewriter, call, replacement); + } + return replacement; +} + +/// Erase an original nondeterministic value after all member writes have moved to replacements. +static void eraseUnusedNondetDef( + Value value, PatternRewriter &rewriter, DenseMap &typedNondetReplacements +) { + if (value.use_empty()) { + if (NonDetOp nondetOp = value.getDefiningOp()) { + typedNondetReplacements.erase(value); + rewriter.eraseOp(nondetOp); + } + } +} + +/// Replace a member-write value and remove an obsolete nondeterministic definition if applicable. +static inline void replaceMemberWriteValue( + MemberWriteOp op, Value originalValue, Value replacement, PatternRewriter &rewriter, + DenseMap &typedNondetReplacements +) { + rewriter.modifyOpInPlace(op, [&op, replacement]() { op.getValMutable().assign(replacement); }); + eraseUnusedNondetDef(originalValue, rewriter, typedNondetReplacements); +} + +} // namespace + +/// Update stale `poly.unifiable_cast` result types after their input has been instantiated. +class UpdateUnifiableCastResultFromInput final : public OpRewritePattern { + ConversionTracker &tracker_; + +public: + /// Construct the unifiable-cast result propagation pattern. + UpdateUnifiableCastResultFromInput(MLIRContext *ctx, ConversionTracker &tracker) + : OpRewritePattern(ctx, 3), tracker_(tracker) {} + + /// Retag a cast result when all remaining uses can observe the same instantiated type. + LogicalResult matchAndRewrite(UnifiableCastOp op, PatternRewriter &rewriter) const override { + Type inputType = op.getInput().getType(); + Value result = op.getResult(); + Type resultType = result.getType(); + if (typesUnify(inputType, resultType) || !allMemberWritesAcceptType(result, inputType)) { + return failure(); + } + bool canUse = Step5_ScalarizeHeterogeneousArrays::canUseScalarizedValueAsType( + resultType, inputType, tracker_, "UpdateUnifiableCastResultFromInput" + ); + if (!canUse) { + return failure(); + } + rewriter.modifyOpInPlace(op, [&result, &inputType]() { result.setType(inputType); }); + LLVM_DEBUG( + llvm::dbgs() << "[UpdateUnifiableCastResultFromInput] updated result type in " << op << '\n' + ); + return success(); + } +}; + /// Update the type of MemberWriteOp value based on updated types from MemberDefOp. class UpdateMemberWriteValFromDef final : public OpRewritePattern { ConversionTracker &tracker_; + mutable DenseMap typedNondetReplacements_; public: + /// Construct the member-write value propagation pattern. UpdateMemberWriteValFromDef(MLIRContext *ctx, ConversionTracker &tracker) : OpRewritePattern(ctx, 3), tracker_(tracker) {} + /// Update a member write operand type from its referenced member definition. LogicalResult matchAndRewrite(MemberWriteOp op, PatternRewriter &rewriter) const override { - return updateMemberRefValFromMemberDef(op, tracker_, rewriter); + SymbolTableCollection tables; + auto def = op.getMemberDefOp(tables); + if (failed(def)) { + return failure(); + } + + Value oldValue = op.getVal(); + Type newValueType = def->get().getType(); + auto cachedReplacement = typedNondetReplacements_.find(oldValue); + if (cachedReplacement != typedNondetReplacements_.end()) { + FailureOr convertedValue = materializeValueForMemberWrite( + op.getLoc(), oldValue, newValueType, rewriter, tracker_, typedNondetReplacements_ + ); + if (failed(convertedValue)) { + return failure(); + } + replaceMemberWriteValue(op, oldValue, *convertedValue, rewriter, typedNondetReplacements_); + LLVM_DEBUG( + llvm::dbgs() << "[UpdateMemberWriteValFromDef] reused materialized value type for " << op + << '\n' + ); + return success(); + } + + Type oldValueType = oldValue.getType(); + if (oldValueType == newValueType || + !tracker_.isLegalConversion(oldValueType, newValueType, "UpdateMemberWriteValFromDef")) { + return failure(); + } + + if (!hasUsesOutside(oldValue, op.getOperation())) { + rewriter.modifyOpInPlace(op, [&oldValue, &newValueType]() { + oldValue.setType(newValueType); + }); + LLVM_DEBUG( + llvm::dbgs() << "[UpdateMemberWriteValFromDef] updated value type in " << op << '\n' + ); + return success(); + } + + FailureOr convertedValue = materializeValueForMemberWrite( + op.getLoc(), oldValue, newValueType, rewriter, tracker_, typedNondetReplacements_ + ); + if (failed(convertedValue)) { + return failure(); + } + replaceMemberWriteValue(op, oldValue, *convertedValue, rewriter, typedNondetReplacements_); + LLVM_DEBUG( + llvm::dbgs() << "[UpdateMemberWriteValFromDef] materialized value type for " << op << '\n' + ); + return success(); } }; +/// Run all type-propagation patterns to a local fixpoint for the current iteration. LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { MLIRContext *ctx = modOp.getContext(); RewritePatternSet patterns(ctx); @@ -2469,22 +5828,24 @@ LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) { // benefit = 6 UpdateInferredResultTypes, // OpTrait::InferTypeOpAdaptor (ReadArrayOp, ExtractArrayOp) // benefit = 3 - UpdateFreeFuncCallOpTypes, // CallOp, targeting non-struct functions - UpdateFuncTypeFromReturn, // FuncDefOp - UpdateNewArrayElemFromWrite, // CreateArrayOp - UpdateArrayElemFromArrRead, // ReadArrayOp - UpdateArrayElemFromArrWrite, // WriteArrayOp - UpdateMemberDefTypeFromWrite, // MemberDefOp - UpdateMemberReadValFromDef, // MemberReadOp - UpdateMemberWriteValFromDef // MemberWriteOp + UpdateFreeFuncCallOpTypes, // CallOp, targeting non-struct functions + UpdateFuncTypeFromReturn, // FuncDefOp + UpdateNewArrayElemFromWrite, // CreateArrayOp + UpdateArrayElemFromArrRead, // ReadArrayOp + UpdateArrayElemFromArrWrite, // WriteArrayOp + UpdateMemberDefTypeFromWrite, // MemberDefOp + UpdateMemberReadValFromDef, // MemberReadOp + UpdateUnifiableCastResultFromInput, // UnifiableCastOp + UpdateMemberWriteValFromDef // MemberWriteOp >(ctx, tracker); return applyAndFoldGreedily(modOp, tracker, std::move(patterns)); } -} // namespace Step5_PropagateTypes +} // namespace Step6_PropagateTypes -namespace Step6_Cleanup { +namespace Step7_Cleanup { +/// Cleanup strategy that preserves symbols reachable from an explicit keep set plus globals. struct FromKeepSet : public CleanupBase { using CleanupBase::CleanupBase; @@ -2593,7 +5954,7 @@ struct FromKeepSet : public CleanupBase { } }; -} // namespace Step6_Cleanup +} // namespace Step7_Cleanup class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { using Base = FlatteningPassBase; @@ -2605,6 +5966,7 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { return m == FlatteningCleanupMode::Unspecified ? FlatteningCleanupMode::Preimage : m; } + /// Run the pass on the current module and signal failure if any flattening step fails. void runOnOperation() override { ModuleOp modOp = getOperation(); if (failed(runOn(modOp))) { @@ -2619,6 +5981,7 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { } } + /// Execute the full flattening pipeline until it reaches a fixpoint or the iteration limit. inline LogicalResult runOn(ModuleOp modOp) { FlatteningCleanupMode effectiveCleanupMode = getEffectiveCleanupMode(); // If the cleanup mode is set to remove anything not reachable from the main struct, do an @@ -2670,6 +6033,11 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { llvm::errs() << DEBUG_TYPE << " failed while instantiating structs in templates\n"; return failure(); } + + LLVM_DEBUG({ + llvm::dbgs() << "[FlatteningPass(count=" << loopCount + << ")] Running step 2: function instantiation\n"; + }); // Instantiate calls to templated functions. if (failed(Step2_InstantiateFunctions::run(modOp, tracker))) { llvm::errs() << DEBUG_TYPE << " failed while instantiating functions in templates\n"; @@ -2678,7 +6046,7 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { LLVM_DEBUG({ llvm::dbgs() << "[FlatteningPass(count=" << loopCount - << ")] Running step 2: loop unrolling\n"; + << ")] Running step 3: loop unrolling\n"; }); // Unroll loops with known iterations. if (failed(Step3_Unroll::run(modOp, tracker))) { @@ -2688,7 +6056,7 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { LLVM_DEBUG({ llvm::dbgs() << "[FlatteningPass(count=" << loopCount - << ")] Running step 3: affine maps instantiation\n"; + << ")] Running step 4: affine maps instantiation\n"; }); // Instantiate affine_map parameters of StructType and ArrayType. if (failed(Step4_InstantiateAffineMaps::run(modOp, tracker))) { @@ -2696,12 +6064,23 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { return failure(); } + // Split static arrays whose affine-map element type instantiates to different concrete + // element types at different indices. LLVM_DEBUG({ llvm::dbgs() << "[FlatteningPass(count=" << loopCount - << ")] Running step 4: type propagation\n"; + << ")] Running step 5: heterogeneous array scalarization\n"; + }); + if (failed(Step5_ScalarizeHeterogeneousArrays::run(modOp, tracker))) { + llvm::errs() << DEBUG_TYPE << " failed while scalarizing heterogeneous arrays\n"; + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "[FlatteningPass(count=" << loopCount + << ")] Running step 6: type propagation\n"; }); // Propagate updated types using the semantics of various ops. - if (failed(Step5_PropagateTypes::run(modOp, tracker))) { + if (failed(Step6_PropagateTypes::run(modOp, tracker))) { llvm::errs() << DEBUG_TYPE << " failed while propagating instantiated types\n"; return failure(); } @@ -2734,10 +6113,10 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { return runPipeline(allocationCleanup, modOp); } - // Perform cleanup according to the 'cleanupMode' option. + /// Perform cleanup according to the effective `cleanupMode` option. LogicalResult cleanupSwitch(ModuleOp modOp, const ConversionTracker &tracker) { FlatteningCleanupMode effectiveCleanupMode = getEffectiveCleanupMode(); - LLVM_DEBUG({ llvm::dbgs() << "[FlatteningPass] Running step 5: cleanup "; }); + LLVM_DEBUG({ llvm::dbgs() << "[FlatteningPass] Running step 7: cleanup "; }); switch (effectiveCleanupMode) { case FlatteningCleanupMode::MainAsRoot: LLVM_DEBUG(llvm::dbgs() << "(main as root mode)\n"); @@ -2748,14 +6127,16 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { case FlatteningCleanupMode::Preimage: LLVM_DEBUG(llvm::dbgs() << "(preimage mode)\n"); return erasePreimageOfInstantiations(modOp, tracker); - case FlatteningCleanupMode::Unspecified: - default: + case FlatteningCleanupMode::Disabled: LLVM_DEBUG(llvm::dbgs() << "(disabled)\n"); return success(); + case FlatteningCleanupMode::Unspecified: + llvm_unreachable("`getEffectiveCleanupMode()` cannot give `Unspecified`"); } + llvm_unreachable("unknown cleanup mode"); } - // Erase parameterized definitions that were replaced with concrete instantiations. + /// Erase parameterized definitions that were replaced with concrete instantiations. LogicalResult erasePreimageOfInstantiations(ModuleOp rootMod, const ConversionTracker &tracker) { // TODO: The names from getInstantiatedDefinitionNames() are NOT guaranteed to be paths from the // "top root" and they also do not indicate a root module so there could be ambiguity. This is a @@ -2790,22 +6171,24 @@ class PassImpl : public llzk::polymorphic::impl::FlatteningPassBase { return res; } + /// Erase cleanup candidates that are unreachable from any concrete definition or global. LogicalResult eraseUnreachableFromConcreteDefinitions(ModuleOp rootMod) { SmallVector roots; rootMod.walk([&roots](Operation *op) { - if (isErasableDefinition(op) && !Step6_Cleanup::FromKeepSet::hasTemplateSymbolBindings(op)) { + if (isErasableDefinition(op) && !Step7_Cleanup::FromKeepSet::hasTemplateSymbolBindings(op)) { roots.push_back(llvm::cast(op)); } }); - Step6_Cleanup::FromKeepSet cleaner( + Step7_Cleanup::FromKeepSet cleaner( rootMod, getAnalysis(), getAnalysis() ); return cleaner.eraseUnreachableFrom(roots); } + /// Erase cleanup candidates that are unreachable from the `llzk.main` struct or globals. LogicalResult eraseUnreachableFromMainStruct(ModuleOp rootMod, bool emitWarning = true) { - Step6_Cleanup::FromKeepSet cleaner( + Step7_Cleanup::FromKeepSet cleaner( rootMod, getAnalysis(), getAnalysis() ); FailureOr> mainOpt = diff --git a/test/Transforms/Flattening/circom_example_2B.llzk b/test/Transforms/Flattening/circom_example_2B.llzk new file mode 100644 index 0000000000..21b6c87512 --- /dev/null +++ b/test/Transforms/Flattening/circom_example_2B.llzk @@ -0,0 +1,427 @@ +// RUN: llzk-opt -llzk-flatten %s 2>&1 | FileCheck --enable-var-scope %s + +#IdxToLen = affine_map<(i)[] -> (5*i+1)> // LLZK should always use dimension identifiers, in parens +!IdxToLenSigArray = !struct.type<@TVarArray::@VarArray<[#IdxToLen, !felt.type]>> + +module attributes {llzk.lang = "circom", llzk.main = !struct.type<@TComputeValue::@ComputeValue<[3]>>} { + poly.template @TGetSum { + poly.param @A : index // instantiations of A are {1, 6, 11, ... 5*(P-1)+1}; P=3 in this example + + struct.def @GetSum { + struct.member @out: !felt.type {llzk.pub} + struct.member @sum: !felt.type + + function.def @compute(%inp: !array.type<@A x !felt.type>) -> !struct.type<@TGetSum::@GetSum<[@A]>> { + %self = struct.new : !struct.type<@TGetSum::@GetSum<[@A]>> + // + %0 = felt.const 0 + %lb = arith.constant 0 : index + %ub = poly.read_const @A : index + %step = arith.constant 1 : index + %sum = scf.for %i = %lb to %ub step %step + iter_args(%cur_sum = %0) -> !felt.type { + %next = array.read %inp[%i] : !array.type<@A x !felt.type>, !felt.type + %new_sum = felt.add %cur_sum, %next + scf.yield %new_sum : !felt.type + } + struct.writem %self[@sum] = %sum : !struct.type<@TGetSum::@GetSum<[@A]>>, !felt.type + struct.writem %self[@out] = %sum : !struct.type<@TGetSum::@GetSum<[@A]>>, !felt.type + function.return %self: !struct.type<@TGetSum::@GetSum<[@A]>> + } + + function.def @constrain(%self: !struct.type<@TGetSum::@GetSum<[@A]>>, %inp: !array.type<@A x !felt.type>) { + %sum = struct.readm %self[@sum] : !struct.type<@TGetSum::@GetSum<[@A]>>, !felt.type + %out = struct.readm %self[@out] : !struct.type<@TGetSum::@GetSum<[@A]>>, !felt.type + constrain.eq %out, %sum : !felt.type + function.return + } + } + } + + // This struct is needed because LLZK doesn't allow array type as an element and the + // type !array.type<@P,#IdxToLen x !felt.type> is not valid because it's not rectangular + // and can't be created via `array.new` outside the loop where the loop induction + // variable does not yet exist. + poly.template @TVarArray { + poly.param @N : index + poly.param @T + + struct.def @VarArray { + struct.member @val: !array.type<@N x !poly.tvar<@T>> {llzk.pub} + + function.def @compute(%inp: !array.type<@N x !poly.tvar<@T>>) -> !struct.type<@TVarArray::@VarArray<[@N, @T]>> { + %self = struct.new : !struct.type<@TVarArray::@VarArray<[@N, @T]>> + struct.writem %self[@val] = %inp : !struct.type<@TVarArray::@VarArray<[@N, @T]>>, !array.type<@N x !poly.tvar<@T>> + function.return %self : !struct.type<@TVarArray::@VarArray<[@N, @T]>> + } + + function.def @constrain(%self: !struct.type<@TVarArray::@VarArray<[@N, @T]>>, %inp: !array.type<@N x !poly.tvar<@T>>) { + function.return + } + } + } + + poly.template @TComputeValue { + poly.param @P : index + + struct.def @ComputeValue { + struct.member @ret: !array.type<@P x !felt.type> {llzk.pub} + struct.member @ws: !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>> + struct.member @arrs: !array.type<@P x !IdxToLenSigArray> + + function.def @compute(%inp: !array.type<@P x !felt.type>) -> !struct.type<@TComputeValue::@ComputeValue<[@P]>> { + %self = struct.new : !struct.type<@TComputeValue::@ComputeValue<[@P]>> + // + %lb = arith.constant 0 : index + %ub = poly.read_const @P : index + %step = arith.constant 1 : index + // for(i = 0; i < @P; i++) { + // len := 5*i+1; + // arr := array.new {len = 5*i+1}; + // for(k = 0; k < len; k++) { + // arr[k] = inp[i] * k; + // } + // ws[i] := GetSum::compute(arr); + // arrs[i] := arr; + // } + %temp_ws = array.new : !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>> + %temp_arrs = array.new : !array.type<@P x !IdxToLenSigArray> + scf.for %i = %lb to %ub step %step { + %inpi = array.read %inp[%i] : !array.type<@P x !felt.type>, !felt.type + // + %arr = array.new{(%i)} : !array.type<#IdxToLen x !felt.type> // lengths are {1, 6, 11, ... 5*P+1} + %len = poly.applymap(%i) #IdxToLen + scf.for %k = %lb to %len step %step { + %t0 = cast.tofelt %k : index + %t1 = felt.mul %inpi, %t0 + array.write %arr[%k] = %t1 : !array.type<#IdxToLen x !felt.type>, !felt.type + } + %wsi = function.call @TGetSum::@GetSum::@compute(%arr){(%i)} : (!array.type<#IdxToLen x !felt.type>) -> !struct.type<@TGetSum::@GetSum<[#IdxToLen]>> + array.write %temp_ws[%i] = %wsi : !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>>, !struct.type<@TGetSum::@GetSum<[#IdxToLen]>> + %arr_wrap = function.call @TVarArray::@VarArray::@compute(%arr){(%i)} : (!array.type<#IdxToLen x !felt.type>) -> !IdxToLenSigArray + array.write %temp_arrs[%i] = %arr_wrap : !array.type<@P x !IdxToLenSigArray>, !IdxToLenSigArray + } + struct.writem %self[@ws] = %temp_ws : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>> + struct.writem %self[@arrs] = %temp_arrs : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !IdxToLenSigArray> + // for(j = 0; j < @P; j++) { + // ret[j] := ws[j].out; + // } + %temp_ret = array.new : !array.type<@P x !felt.type> + scf.for %j = %lb to %ub step %step { + %wsj = array.read %temp_ws[%j] : !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>>, !struct.type<@TGetSum::@GetSum<[#IdxToLen]>> + %wsjout = struct.readm %wsj[@out] : !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>, !felt.type + array.write %temp_ret[%j] = %wsjout : !array.type<@P x !felt.type>, !felt.type + } + struct.writem %self[@ret] = %temp_ret : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !felt.type> + // + function.return %self: !struct.type<@TComputeValue::@ComputeValue<[@P]>> + } + + function.def @constrain(%self: !struct.type<@TComputeValue::@ComputeValue<[@P]>>, %inp: !array.type<@P x !felt.type>) { + %lb = arith.constant 0 : index + %ub = poly.read_const @P : index + %step = arith.constant 1 : index + // for(i = 0; i < @P; i++) { + // ws[i].constrain(arrs[i]); + // } + %temp_ws = struct.readm %self[@ws] : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>> + %temp_arrs = struct.readm %self[@arrs] : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !IdxToLenSigArray> + scf.for %i = %lb to %ub step %step { + %wsi = array.read %temp_ws[%i] : !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>>, !struct.type<@TGetSum::@GetSum<[#IdxToLen]>> + %inpi = array.read %inp[%i] : !array.type<@P x !felt.type>, !felt.type + %arr_wrap = array.read %temp_arrs[%i] : !array.type<@P x !IdxToLenSigArray>, !IdxToLenSigArray + %arr = struct.readm %arr_wrap[@val] : !IdxToLenSigArray, !array.type<#IdxToLen x !felt.type> + function.call @TGetSum::@GetSum::@constrain(%wsi, %arr) : (!struct.type<@TGetSum::@GetSum<[#IdxToLen]>>, !array.type<#IdxToLen x !felt.type>) -> () + } + // for(j = 0; j < @P; j++) { + // emit ret[j] = ws[j].out; + // } + %temp_ret = struct.readm %self[@ret] : !struct.type<@TComputeValue::@ComputeValue<[@P]>>, !array.type<@P x !felt.type> + scf.for %j = %lb to %ub step %step { + %retj = array.read %temp_ret[%j] : !array.type<@P x !felt.type>, !felt.type + %wsj = array.read %temp_ws[%j] : !array.type<@P x !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>>, !struct.type<@TGetSum::@GetSum<[#IdxToLen]>> + %wsjout = struct.readm %wsj[@out] : !struct.type<@TGetSum::@GetSum<[#IdxToLen]>>, !felt.type + constrain.eq %retj, %wsjout : !felt.type + } + + function.return + } + } + } +} + +// CHECK-LABEL: module attributes {llzk.lang = "circom", llzk.main = !struct.type<@TComputeValue_3_ComputeValue>} { +// CHECK-NEXT: struct.def @TGetSum_1_GetSum { +// CHECK-NEXT: struct.member @out : !felt.type {llzk.pub} +// CHECK-NEXT: struct.member @sum : !felt.type +// CHECK-NEXT: function.def @compute(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<1 x !felt.type>) -> !struct.type<@TGetSum_1_GetSum> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = felt.const 0 +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = struct.new : <@TGetSum_1_GetSum> +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_0]]{{\[}}%[[VAL_1]]] : <1 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_4]], %[[VAL_2]] : !felt.type, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_3]][@sum] = %[[VAL_5]] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_3]][@out] = %[[VAL_5]] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: function.return %[[VAL_3]] : !struct.type<@TGetSum_1_GetSum> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@TGetSum_1_GetSum>, %[[VAL_7:[0-9a-zA-Z_\.]+]]: !array.type<1 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_6]][@sum] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_6]][@out] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_9]], %[[VAL_8]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TGetSum_6_GetSum { +// CHECK-NEXT: struct.member @out : !felt.type {llzk.pub} +// CHECK-NEXT: struct.member @sum : !felt.type +// CHECK-NEXT: function.def @compute(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !array.type<6 x !felt.type>) -> !struct.type<@TGetSum_6_GetSum> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = arith.constant 4 : index +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = arith.constant 3 : index +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index +// CHECK-NEXT: %[[VAL_15:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index +// CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_17:[0-9a-zA-Z_\.]+]] = felt.const 0 +// CHECK-NEXT: %[[VAL_18:[0-9a-zA-Z_\.]+]] = struct.new : <@TGetSum_6_GetSum> +// CHECK-NEXT: %[[VAL_19:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_16]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_20:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_19]], %[[VAL_17]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_21:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_15]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_22:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_20]], %[[VAL_21]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_23:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_14]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_24:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_22]], %[[VAL_23]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_25:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_13]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_26:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_24]], %[[VAL_25]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_27:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_12]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_28:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_26]], %[[VAL_27]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_29:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_10]]{{\[}}%[[VAL_11]]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_30:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_28]], %[[VAL_29]] : !felt.type, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_18]][@sum] = %[[VAL_30]] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_18]][@out] = %[[VAL_30]] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: function.return %[[VAL_18]] : !struct.type<@TGetSum_6_GetSum> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_31:[0-9a-zA-Z_\.]+]]: !struct.type<@TGetSum_6_GetSum>, %[[VAL_32:[0-9a-zA-Z_\.]+]]: !array.type<6 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_33:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_31]][@sum] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: %[[VAL_34:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_31]][@out] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_34]], %[[VAL_33]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TGetSum_11_GetSum { +// CHECK-NEXT: struct.member @out : !felt.type {llzk.pub} +// CHECK-NEXT: struct.member @sum : !felt.type +// CHECK-NEXT: function.def @compute(%[[VAL_35:[0-9a-zA-Z_\.]+]]: !array.type<11 x !felt.type>) -> !struct.type<@TGetSum_11_GetSum> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_36:[0-9a-zA-Z_\.]+]] = arith.constant 10 : index +// CHECK-NEXT: %[[VAL_37:[0-9a-zA-Z_\.]+]] = arith.constant 9 : index +// CHECK-NEXT: %[[VAL_38:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_39:[0-9a-zA-Z_\.]+]] = arith.constant 7 : index +// CHECK-NEXT: %[[VAL_40:[0-9a-zA-Z_\.]+]] = arith.constant 6 : index +// CHECK-NEXT: %[[VAL_41:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index +// CHECK-NEXT: %[[VAL_42:[0-9a-zA-Z_\.]+]] = arith.constant 4 : index +// CHECK-NEXT: %[[VAL_43:[0-9a-zA-Z_\.]+]] = arith.constant 3 : index +// CHECK-NEXT: %[[VAL_44:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index +// CHECK-NEXT: %[[VAL_45:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index +// CHECK-NEXT: %[[VAL_46:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_47:[0-9a-zA-Z_\.]+]] = felt.const 0 +// CHECK-NEXT: %[[VAL_48:[0-9a-zA-Z_\.]+]] = struct.new : <@TGetSum_11_GetSum> +// CHECK-NEXT: %[[VAL_49:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_46]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_50:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_49]], %[[VAL_47]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_51:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_45]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_52:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_50]], %[[VAL_51]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_53:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_44]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_54:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_52]], %[[VAL_53]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_55:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_43]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_56:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_54]], %[[VAL_55]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_57:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_42]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_58:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_56]], %[[VAL_57]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_59:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_41]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_60:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_58]], %[[VAL_59]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_61:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_40]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_62:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_60]], %[[VAL_61]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_63:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_39]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_64:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_62]], %[[VAL_63]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_65:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_38]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_66:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_64]], %[[VAL_65]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_67:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_37]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_68:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_66]], %[[VAL_67]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_69:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_35]]{{\[}}%[[VAL_36]]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_70:[0-9a-zA-Z_\.]+]] = felt.add %[[VAL_68]], %[[VAL_69]] : !felt.type, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_48]][@sum] = %[[VAL_70]] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_48]][@out] = %[[VAL_70]] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: function.return %[[VAL_48]] : !struct.type<@TGetSum_11_GetSum> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_71:[0-9a-zA-Z_\.]+]]: !struct.type<@TGetSum_11_GetSum>, %[[VAL_72:[0-9a-zA-Z_\.]+]]: !array.type<11 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_73:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_71]][@sum] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: %[[VAL_74:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_71]][@out] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_74]], %[[VAL_73]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TVarArray_1_f_VarArray { +// CHECK-NEXT: struct.member @val : !array.type<1 x !felt.type> {llzk.pub} +// CHECK-NEXT: function.def @compute(%[[VAL_75:[0-9a-zA-Z_\.]+]]: !array.type<1 x !felt.type>) -> !struct.type<@TVarArray_1_f_VarArray> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_76:[0-9a-zA-Z_\.]+]] = struct.new : <@TVarArray_1_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_76]][@val] = %[[VAL_75]] : <@TVarArray_1_f_VarArray>, !array.type<1 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_76]] : !struct.type<@TVarArray_1_f_VarArray> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_77:[0-9a-zA-Z_\.]+]]: !struct.type<@TVarArray_1_f_VarArray>, %[[VAL_78:[0-9a-zA-Z_\.]+]]: !array.type<1 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TVarArray_6_f_VarArray { +// CHECK-NEXT: struct.member @val : !array.type<6 x !felt.type> {llzk.pub} +// CHECK-NEXT: function.def @compute(%[[VAL_79:[0-9a-zA-Z_\.]+]]: !array.type<6 x !felt.type>) -> !struct.type<@TVarArray_6_f_VarArray> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_80:[0-9a-zA-Z_\.]+]] = struct.new : <@TVarArray_6_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_80]][@val] = %[[VAL_79]] : <@TVarArray_6_f_VarArray>, !array.type<6 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_80]] : !struct.type<@TVarArray_6_f_VarArray> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_81:[0-9a-zA-Z_\.]+]]: !struct.type<@TVarArray_6_f_VarArray>, %[[VAL_82:[0-9a-zA-Z_\.]+]]: !array.type<6 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TVarArray_11_f_VarArray { +// CHECK-NEXT: struct.member @val : !array.type<11 x !felt.type> {llzk.pub} +// CHECK-NEXT: function.def @compute(%[[VAL_83:[0-9a-zA-Z_\.]+]]: !array.type<11 x !felt.type>) -> !struct.type<@TVarArray_11_f_VarArray> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_84:[0-9a-zA-Z_\.]+]] = struct.new : <@TVarArray_11_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_84]][@val] = %[[VAL_83]] : <@TVarArray_11_f_VarArray>, !array.type<11 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_84]] : !struct.type<@TVarArray_11_f_VarArray> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_85:[0-9a-zA-Z_\.]+]]: !struct.type<@TVarArray_11_f_VarArray>, %[[VAL_86:[0-9a-zA-Z_\.]+]]: !array.type<11 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TComputeValue_3_ComputeValue { +// CHECK-NEXT: struct.member @ret : !array.type<3 x !felt.type> {llzk.pub} +// CHECK-NEXT: struct.member @ws_0 : !struct.type<@TGetSum_1_GetSum> +// CHECK-NEXT: struct.member @ws_1 : !struct.type<@TGetSum_6_GetSum> +// CHECK-NEXT: struct.member @ws_2 : !struct.type<@TGetSum_11_GetSum> +// CHECK-NEXT: struct.member @arrs_3 : !struct.type<@TVarArray_1_f_VarArray> +// CHECK-NEXT: struct.member @arrs_4 : !struct.type<@TVarArray_6_f_VarArray> +// CHECK-NEXT: struct.member @arrs_5 : !struct.type<@TVarArray_11_f_VarArray> +// CHECK-NEXT: function.def @compute(%[[VAL_87:[0-9a-zA-Z_\.]+]]: !array.type<3 x !felt.type>) -> !struct.type<@TComputeValue_3_ComputeValue> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_88:[0-9a-zA-Z_\.]+]] = arith.constant 10 : index +// CHECK-NEXT: %[[VAL_89:[0-9a-zA-Z_\.]+]] = arith.constant 9 : index +// CHECK-NEXT: %[[VAL_90:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_91:[0-9a-zA-Z_\.]+]] = arith.constant 7 : index +// CHECK-NEXT: %[[VAL_92:[0-9a-zA-Z_\.]+]] = arith.constant 6 : index +// CHECK-NEXT: %[[VAL_93:[0-9a-zA-Z_\.]+]] = arith.constant 5 : index +// CHECK-NEXT: %[[VAL_94:[0-9a-zA-Z_\.]+]] = arith.constant 4 : index +// CHECK-NEXT: %[[VAL_95:[0-9a-zA-Z_\.]+]] = arith.constant 3 : index +// CHECK-NEXT: %[[VAL_96:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index +// CHECK-NEXT: %[[VAL_97:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index +// CHECK-NEXT: %[[VAL_98:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_99:[0-9a-zA-Z_\.]+]] = struct.new : <@TComputeValue_3_ComputeValue> +// CHECK-NEXT: %[[VAL_100:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_87]]{{\[}}%[[VAL_98]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_101:[0-9a-zA-Z_\.]+]] = array.new : <1 x !felt.type> +// CHECK-NEXT: %[[VAL_102:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_98]] : index, !felt.type +// CHECK-NEXT: %[[VAL_103:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_100]], %[[VAL_102]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_101]]{{\[}}%[[VAL_98]]] = %[[VAL_103]] : <1 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_104:[0-9a-zA-Z_\.]+]] = function.call @TGetSum_1_GetSum::@compute(%[[VAL_101]]) : (!array.type<1 x !felt.type>) -> !struct.type<@TGetSum_1_GetSum> +// CHECK-NEXT: %[[VAL_105:[0-9a-zA-Z_\.]+]] = function.call @TVarArray_1_f_VarArray::@compute(%[[VAL_101]]) : (!array.type<1 x !felt.type>) -> !struct.type<@TVarArray_1_f_VarArray> +// CHECK-NEXT: %[[VAL_106:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_87]]{{\[}}%[[VAL_97]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_107:[0-9a-zA-Z_\.]+]] = array.new : <6 x !felt.type> +// CHECK-NEXT: %[[VAL_108:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_98]] : index, !felt.type +// CHECK-NEXT: %[[VAL_109:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_108]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_98]]] = %[[VAL_109]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_110:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_97]] : index, !felt.type +// CHECK-NEXT: %[[VAL_111:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_110]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_97]]] = %[[VAL_111]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_112:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_96]] : index, !felt.type +// CHECK-NEXT: %[[VAL_113:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_112]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_96]]] = %[[VAL_113]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_114:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_95]] : index, !felt.type +// CHECK-NEXT: %[[VAL_115:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_114]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_95]]] = %[[VAL_115]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_116:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_94]] : index, !felt.type +// CHECK-NEXT: %[[VAL_117:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_116]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_94]]] = %[[VAL_117]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_118:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_93]] : index, !felt.type +// CHECK-NEXT: %[[VAL_119:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_106]], %[[VAL_118]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_107]]{{\[}}%[[VAL_93]]] = %[[VAL_119]] : <6 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_120:[0-9a-zA-Z_\.]+]] = function.call @TGetSum_6_GetSum::@compute(%[[VAL_107]]) : (!array.type<6 x !felt.type>) -> !struct.type<@TGetSum_6_GetSum> +// CHECK-NEXT: %[[VAL_121:[0-9a-zA-Z_\.]+]] = function.call @TVarArray_6_f_VarArray::@compute(%[[VAL_107]]) : (!array.type<6 x !felt.type>) -> !struct.type<@TVarArray_6_f_VarArray> +// CHECK-NEXT: %[[VAL_122:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_87]]{{\[}}%[[VAL_96]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_123:[0-9a-zA-Z_\.]+]] = array.new : <11 x !felt.type> +// CHECK-NEXT: %[[VAL_124:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_98]] : index, !felt.type +// CHECK-NEXT: %[[VAL_125:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_124]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_98]]] = %[[VAL_125]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_126:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_97]] : index, !felt.type +// CHECK-NEXT: %[[VAL_127:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_126]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_97]]] = %[[VAL_127]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_128:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_96]] : index, !felt.type +// CHECK-NEXT: %[[VAL_129:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_128]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_96]]] = %[[VAL_129]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_130:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_95]] : index, !felt.type +// CHECK-NEXT: %[[VAL_131:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_130]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_95]]] = %[[VAL_131]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_132:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_94]] : index, !felt.type +// CHECK-NEXT: %[[VAL_133:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_132]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_94]]] = %[[VAL_133]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_134:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_93]] : index, !felt.type +// CHECK-NEXT: %[[VAL_135:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_134]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_93]]] = %[[VAL_135]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_136:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_92]] : index, !felt.type +// CHECK-NEXT: %[[VAL_137:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_136]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_92]]] = %[[VAL_137]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_138:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_91]] : index, !felt.type +// CHECK-NEXT: %[[VAL_139:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_138]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_91]]] = %[[VAL_139]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_140:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_90]] : index, !felt.type +// CHECK-NEXT: %[[VAL_141:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_140]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_90]]] = %[[VAL_141]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_142:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_89]] : index, !felt.type +// CHECK-NEXT: %[[VAL_143:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_142]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_89]]] = %[[VAL_143]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_144:[0-9a-zA-Z_\.]+]] = cast.tofelt %[[VAL_88]] : index, !felt.type +// CHECK-NEXT: %[[VAL_145:[0-9a-zA-Z_\.]+]] = felt.mul %[[VAL_122]], %[[VAL_144]] : !felt.type, !felt.type +// CHECK-NEXT: array.write %[[VAL_123]]{{\[}}%[[VAL_88]]] = %[[VAL_145]] : <11 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_146:[0-9a-zA-Z_\.]+]] = function.call @TGetSum_11_GetSum::@compute(%[[VAL_123]]) : (!array.type<11 x !felt.type>) -> !struct.type<@TGetSum_11_GetSum> +// CHECK-NEXT: %[[VAL_147:[0-9a-zA-Z_\.]+]] = function.call @TVarArray_11_f_VarArray::@compute(%[[VAL_123]]) : (!array.type<11 x !felt.type>) -> !struct.type<@TVarArray_11_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_99]][@ws_0] = %[[VAL_104]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_1_GetSum> +// CHECK-NEXT: struct.writem %[[VAL_99]][@ws_1] = %[[VAL_120]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_6_GetSum> +// CHECK-NEXT: struct.writem %[[VAL_99]][@ws_2] = %[[VAL_146]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_11_GetSum> +// CHECK-NEXT: struct.writem %[[VAL_99]][@arrs_3] = %[[VAL_105]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_1_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_99]][@arrs_4] = %[[VAL_121]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_6_f_VarArray> +// CHECK-NEXT: struct.writem %[[VAL_99]][@arrs_5] = %[[VAL_147]] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_11_f_VarArray> +// CHECK-NEXT: %[[VAL_148:[0-9a-zA-Z_\.]+]] = array.new : <3 x !felt.type> +// CHECK-NEXT: %[[VAL_149:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_104]][@out] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: array.write %[[VAL_148]]{{\[}}%[[VAL_98]]] = %[[VAL_149]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_150:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_120]][@out] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: array.write %[[VAL_148]]{{\[}}%[[VAL_97]]] = %[[VAL_150]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_151:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_146]][@out] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: array.write %[[VAL_148]]{{\[}}%[[VAL_96]]] = %[[VAL_151]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: struct.writem %[[VAL_99]][@ret] = %[[VAL_148]] : <@TComputeValue_3_ComputeValue>, !array.type<3 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_99]] : !struct.type<@TComputeValue_3_ComputeValue> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_152:[0-9a-zA-Z_\.]+]]: !struct.type<@TComputeValue_3_ComputeValue>, %[[VAL_153:[0-9a-zA-Z_\.]+]]: !array.type<3 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_154:[0-9a-zA-Z_\.]+]] = arith.constant 2 : index +// CHECK-NEXT: %[[VAL_155:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_156:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index +// CHECK-NEXT: %[[VAL_157:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@ws_2] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_11_GetSum> +// CHECK-NEXT: %[[VAL_158:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@ws_1] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_6_GetSum> +// CHECK-NEXT: %[[VAL_159:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@ws_0] : <@TComputeValue_3_ComputeValue>, !struct.type<@TGetSum_1_GetSum> +// CHECK-NEXT: %[[VAL_160:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@arrs_5] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_11_f_VarArray> +// CHECK-NEXT: %[[VAL_161:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@arrs_4] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_6_f_VarArray> +// CHECK-NEXT: %[[VAL_162:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@arrs_3] : <@TComputeValue_3_ComputeValue>, !struct.type<@TVarArray_1_f_VarArray> +// CHECK-NEXT: %[[VAL_163:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_153]]{{\[}}%[[VAL_155]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_164:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_162]][@val] : <@TVarArray_1_f_VarArray>, !array.type<1 x !felt.type> +// CHECK-NEXT: function.call @TGetSum_1_GetSum::@constrain(%[[VAL_159]], %[[VAL_164]]) : (!struct.type<@TGetSum_1_GetSum>, !array.type<1 x !felt.type>) -> () +// CHECK-NEXT: %[[VAL_165:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_153]]{{\[}}%[[VAL_156]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_166:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_161]][@val] : <@TVarArray_6_f_VarArray>, !array.type<6 x !felt.type> +// CHECK-NEXT: function.call @TGetSum_6_GetSum::@constrain(%[[VAL_158]], %[[VAL_166]]) : (!struct.type<@TGetSum_6_GetSum>, !array.type<6 x !felt.type>) -> () +// CHECK-NEXT: %[[VAL_167:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_153]]{{\[}}%[[VAL_154]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_168:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_160]][@val] : <@TVarArray_11_f_VarArray>, !array.type<11 x !felt.type> +// CHECK-NEXT: function.call @TGetSum_11_GetSum::@constrain(%[[VAL_157]], %[[VAL_168]]) : (!struct.type<@TGetSum_11_GetSum>, !array.type<11 x !felt.type>) -> () +// CHECK-NEXT: %[[VAL_169:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_152]][@ret] : <@TComputeValue_3_ComputeValue>, !array.type<3 x !felt.type> +// CHECK-NEXT: %[[VAL_170:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_169]]{{\[}}%[[VAL_155]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_171:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_159]][@out] : <@TGetSum_1_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_170]], %[[VAL_171]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_172:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_169]]{{\[}}%[[VAL_156]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_173:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_158]][@out] : <@TGetSum_6_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_172]], %[[VAL_173]] : !felt.type, !felt.type +// CHECK-NEXT: %[[VAL_174:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_169]]{{\[}}%[[VAL_154]]] : <3 x !felt.type>, !felt.type +// CHECK-NEXT: %[[VAL_175:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_157]][@out] : <@TGetSum_11_GetSum>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_174]], %[[VAL_175]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/expandable_read_consumer_preflight_fail.llzk b/test/Transforms/Flattening/expandable_read_consumer_preflight_fail.llzk new file mode 100644 index 0000000000..0bbf20c397 --- /dev/null +++ b/test/Transforms/Flattening/expandable_read_consumer_preflight_fail.llzk @@ -0,0 +1,57 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +// A non-candidate whole-array write must validate generic static-read consumers during +// preflight. Otherwise the candidate write can be rewritten before this read fails validation. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %candidate = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %candidate[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %candidate[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %candidate : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a function call}} + %read0 = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{call argument requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + function.call @consume1(%read0) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/instantiate_column_field_fail.llzk b/test/Transforms/Flattening/instantiate_column_field_fail.llzk index da2bbd1e40..08c3ad700f 100644 --- a/test/Transforms/Flattening/instantiate_column_field_fail.llzk +++ b/test/Transforms/Flattening/instantiate_column_field_fail.llzk @@ -84,7 +84,7 @@ module attributes {llzk.lang} { } function.def @main(%data: !struct.type<@Data>) -> !felt.type<"bn128"> { - // expected-error@+1 {{failure while creating instantiated function 'OffsetTemplate_f<35:5:bn128>_read'}} + // expected-error@+1 {{failure while creating instantiated function "OffsetTemplate_f<35:5:bn128>_read"}} %value = function.call @OffsetTemplate::@read<[#felt>]>(%data) : (!struct.type<@Data>) -> !felt.type<"bn128"> function.return %value : !felt.type<"bn128"> } diff --git a/test/Transforms/Flattening/instantiate_func_full_cache_across_iterations.llzk b/test/Transforms/Flattening/instantiate_func_full_cache_across_iterations.llzk new file mode 100644 index 0000000000..85db79135e --- /dev/null +++ b/test/Transforms/Flattening/instantiate_func_full_cache_across_iterations.llzk @@ -0,0 +1,101 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +// Regression for full function-instantiation cache lifetime across flattening iterations. +// One call to @F::@f is concretely @N=4 before unrolling, while another becomes @N=4 only +// after Step 3 unroll + Step 4 affine-map folding. Step 2 must reuse the same full clone. + +#len4 = affine_map<(i) -> (4)> +module attributes {llzk.lang} { + poly.template @F { + poly.param @N : index + + function.def @f(%inp: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { + function.return %inp : !array.type<@N x !felt.type> + } + } + + poly.template @TWrap { + poly.param @N : index + + struct.def @Wrap { + function.def @compute(%inp: !array.type<@N x !felt.type>) -> !struct.type<@TWrap::@Wrap<[@N]>> { + %tmp = function.call @F::@f(%inp) : (!array.type<@N x !felt.type>) -> (!array.type<@N x !felt.type>) + %self = struct.new : !struct.type<@TWrap::@Wrap<[@N]>> + function.return %self : !struct.type<@TWrap::@Wrap<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TWrap::@Wrap<[@N]>>, %inp: !array.type<@N x !felt.type>) { + function.return + } + } + } + + struct.def @Main { + function.def @compute( + %a: !array.type<4 x !felt.type>, + %b: !array.type<4 x !felt.type> + ) -> !struct.type<@Main> { + // First call is concrete in the first Step 2 iteration. + %early = function.call @F::@f(%a) : (!array.type<4 x !felt.type>) -> (!array.type<4 x !felt.type>) + + // This call's result type starts with an affine-map parameter and becomes concrete only + // after Step 4 folds map operands. In the next iteration, Step 1 instantiates @TWrap::@Wrap + // and introduces a second concrete @F::@f call in the cloned body; that call must reuse + // @F_4_f instead of creating a uniquified duplicate. + %c0 = arith.constant 0 : index + %lateStruct = function.call @TWrap::@Wrap::@compute(%b){(%c0)} + : (!array.type<4 x !felt.type>) -> (!struct.type<@TWrap::@Wrap<[#len4]>>) + + %self = struct.new : !struct.type<@Main> + function.return %self : !struct.type<@Main> + } + + function.def @constrain( + %self: !struct.type<@Main>, + %a: !array.type<4 x !felt.type>, + %b: !array.type<4 x !felt.type> + ) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: function.def @F_4_f( +// CHECK-SAME: %[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type> +// CHECK-SAME: ) -> !array.type<4 x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_0]] : !array.type<4 x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TWrap_4_Wrap { +// CHECK-NEXT: function.def @compute( +// CHECK-SAME: %[[VAL_1:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type> +// CHECK-SAME: ) -> !struct.type<@TWrap_4_Wrap> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = function.call @F_4_f(%[[VAL_1]]) : (!array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = struct.new : <@TWrap_4_Wrap> +// CHECK-NEXT: function.return %[[VAL_3]] : !struct.type<@TWrap_4_Wrap> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain( +// CHECK-SAME: %[[VAL_4:[0-9a-zA-Z_\.]+]]: !struct.type<@TWrap_4_Wrap>, +// CHECK-SAME: %[[VAL_5:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type> +// CHECK-SAME: ) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Main { +// CHECK-NEXT: function.def @compute( +// CHECK-SAME: %[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>, +// CHECK-SAME: %[[VAL_7:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type> +// CHECK-SAME: ) -> !struct.type<@Main> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @F_4_f(%[[VAL_6]]) : (!array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TWrap_4_Wrap::@compute(%[[VAL_7]]) : (!array.type<4 x !felt.type>) -> !struct.type<@TWrap_4_Wrap> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@Main> +// CHECK-NEXT: function.return %[[VAL_10]] : !struct.type<@Main> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain( +// CHECK-SAME: %[[VAL_11:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, +// CHECK-SAME: %[[VAL_12:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>, +// CHECK-SAME: %[[VAL_13:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type> +// CHECK-SAME: ) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/instantiate_func_name_collision.llzk b/test/Transforms/Flattening/instantiate_func_name_collision.llzk new file mode 100644 index 0000000000..c8f7a25c4d --- /dev/null +++ b/test/Transforms/Flattening/instantiate_func_name_collision.llzk @@ -0,0 +1,52 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +module attributes {llzk.lang} { + // Deliberately collides with the generated full-instantiation name for + // @template_collision::@f<[4]>. Flattening must create a uniquified clone + // rather than redirecting the call to this unrelated function. + function.def @template_collision_4_f(%inp: !array.type<4 x !felt.type>) -> !felt.type { + %c0 = arith.constant 0 : index + %r = array.read %inp[%c0] : !array.type<4 x !felt.type>, !felt.type + function.return %r : !felt.type + } + + poly.template @template_collision { + poly.param @N : index + + function.def @f(%inp: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { + function.return %inp : !array.type<@N x !felt.type> + } + } + + struct.def @Main { + function.def @compute(%inp: !array.type<4 x !felt.type>) -> !struct.type<@Main> { + %self = struct.new : !struct.type<@Main> + function.call @template_collision::@f(%inp) : (!array.type<4 x !felt.type>) -> (!array.type<4 x !felt.type>) + function.return %self : !struct.type<@Main> + } + + function.def @constrain(%self: !struct.type<@Main>, %inp: !array.type<4 x !felt.type>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: function.def @template_collision_4_f(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_0]]{{\[}}%[[VAL_1]]] : <4 x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_2]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: function.def @template_collision_4_f_0(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_3]] : !array.type<4 x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Main { +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Main> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @template_collision_4_f_0(%[[VAL_4]]) : (!array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Main> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/instantiate_func_nested_module_cache.llzk b/test/Transforms/Flattening/instantiate_func_nested_module_cache.llzk new file mode 100644 index 0000000000..16abfc4935 --- /dev/null +++ b/test/Transforms/Flattening/instantiate_func_nested_module_cache.llzk @@ -0,0 +1,94 @@ +// RUN: llzk-opt -verify-diagnostics --pass-pipeline='builtin.module(llzk-flatten{cleanup=disabled})' %s | FileCheck --enable-var-scope %s + +// Isolate function-instantiation cache behavior from cleanup pruning. Each sibling root defines +// the same relative @T::@f symbol and instantiates the same argument list. + +module attributes {llzk.lang} { + module @Left attributes {llzk.lang} { + poly.template @T { + poly.param @N : index + + function.def @f(%inp: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { + function.return %inp : !array.type<@N x !felt.type> + } + } + + struct.def @Main { + function.def @compute(%inp: !array.type<4 x !felt.type>) -> !struct.type<@Main> { + %self = struct.new : !struct.type<@Main> + %out = function.call @T::@f(%inp) : (!array.type<4 x !felt.type>) -> (!array.type<4 x !felt.type>) + function.return %self : !struct.type<@Main> + } + + function.def @constrain(%self: !struct.type<@Main>, %inp: !array.type<4 x !felt.type>) { + function.return + } + } + } + + module @Right attributes {llzk.lang} { + poly.template @T { + poly.param @N : index + + function.def @f(%inp: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { + function.return %inp : !array.type<@N x !felt.type> + } + } + + struct.def @Main { + function.def @compute(%inp: !array.type<4 x !felt.type>) -> !struct.type<@Main> { + %self = struct.new : !struct.type<@Main> + %out = function.call @T::@f(%inp) : (!array.type<4 x !felt.type>) -> (!array.type<4 x !felt.type>) + function.return %self : !struct.type<@Main> + } + + function.def @constrain(%self: !struct.type<@Main>, %inp: !array.type<4 x !felt.type>) { + function.return + } + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: module @Left attributes {llzk.lang} { +// CHECK-NEXT: function.def @T_4_f(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_0]] : !array.type<4 x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @T { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: function.def @f(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_1]] : !array.type<@N x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Main { +// CHECK-NEXT: function.def @compute(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = struct.new : <@Main> +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = function.call @T_4_f(%[[VAL_2]]) : (!array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_3]] : !struct.type<@Main> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: module @Right attributes {llzk.lang} { +// CHECK-NEXT: function.def @T_4_f(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_7]] : !array.type<4 x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @T { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: function.def @f(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<@N x !felt.type>) -> !array.type<@N x !felt.type> { +// CHECK-NEXT: function.return %[[VAL_8]] : !array.type<@N x !felt.type> +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Main { +// CHECK-NEXT: function.def @compute(%[[VAL_9:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !struct.type<@Main> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.new : <@Main> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @T_4_f(%[[VAL_9]]) : (!array.type<4 x !felt.type>) -> !array.type<4 x !felt.type> +// CHECK-NEXT: function.return %[[VAL_10]] : !struct.type<@Main> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@Main>, %[[VAL_13:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/instantiate_funcs_fail.llzk b/test/Transforms/Flattening/instantiate_funcs_fail.llzk index b4e5b97c21..27976bbe9f 100644 --- a/test/Transforms/Flattening/instantiate_funcs_fail.llzk +++ b/test/Transforms/Flattening/instantiate_funcs_fail.llzk @@ -26,7 +26,7 @@ module attributes {llzk.lang} { struct.def @Main { function.def @compute(%z: !array.type<8,9 x i1>) -> !struct.type<@Main> { %self = struct.new : !struct.type<@Main> - // expected-error@+1 {{failure while creating instantiated function 'template_a2_8_9_b_f'}} + // expected-error@+1 {{failure while creating instantiated function "template_a2_8_9_b_f"}} function.call @template_a2::@f(%z) : (!array.type<8,9 x i1>) -> i1 function.return %self : !struct.type<@Main> } @@ -64,7 +64,7 @@ module attributes {llzk.lang} { struct.def @Main { function.def @compute(%z: !array.type<8,9 x i1>) -> !struct.type<@Main> { %self = struct.new : !struct.type<@Main> - // expected-error@+1 {{failure while creating instantiated function 'template_b2_8_9_b_f'}} + // expected-error@+1 {{failure while creating instantiated function "template_b2_8_9_b_f"}} function.call @template_b2::@f<[8,9,i1]>(%z) : (!array.type<8,9 x i1>) -> i1 function.return %self : !struct.type<@Main> } @@ -106,7 +106,7 @@ module attributes {llzk.lang} { struct.def @Main { function.def @compute(%z: !array.type<8,9 x i1>) -> !struct.type<@Main> { %self = struct.new : !struct.type<@Main> - // expected-error@+1 {{failure while creating instantiated function 'template_c2_8_9_b_f'}} + // expected-error@+1 {{failure while creating instantiated function "template_c2_8_9_b_f"}} function.call @template_c2::@f(%z) : (!array.type<8,9 x i1>) -> i1 function.return %self : !struct.type<@Main> } diff --git a/test/Transforms/Flattening/instantiate_funcs_pass.llzk b/test/Transforms/Flattening/instantiate_funcs_pass.llzk index 89468d7776..0e386e473f 100644 --- a/test/Transforms/Flattening/instantiate_funcs_pass.llzk +++ b/test/Transforms/Flattening/instantiate_funcs_pass.llzk @@ -67,6 +67,55 @@ module attributes {llzk.lang} { // CHECK-NEXT: } // ----- +// Nested type parameters are scoped to their own template, even when they have the same +// names as concrete parameters in the enclosing template. +module attributes {llzk.lang} { + poly.template @NestedScopeInner { + poly.param @T : !poly.tvar<@T> + poly.param @U : !poly.tvar<@U> + + function.def @f(%arg: !poly.tvar<@U>) -> !poly.tvar<@U> { + %cast = poly.unifiable_cast %arg : (!poly.tvar<@U>) -> !poly.tvar<@T> + function.return %arg : !poly.tvar<@U> + } + } + + poly.template @NestedScopeOuter { + poly.param @T : !poly.tvar<@T> + poly.param @U : !poly.tvar<@U> + + function.def @g(%arg: !poly.tvar<@U>) -> !poly.tvar<@U> { + %call = function.call @NestedScopeInner::@f<[?, @U]>(%arg) : (!poly.tvar<@U>) -> !poly.tvar<@U> + function.return %call : !poly.tvar<@U> + } + } + + poly.template @NestedScopeCaller { + poly.param @U : !poly.tvar<@U> + + function.def @call(%arg: !poly.tvar<@U>) -> !poly.tvar<@U> { + %call = function.call @NestedScopeOuter::@g<[index, ?]>(%arg) : (!poly.tvar<@U>) -> !poly.tvar<@U> + function.return %call : !poly.tvar<@U> + } + } +} +// CHECK-LABEL: poly.template @"NestedScopeOuter_i_\1A" { +// CHECK-NEXT: poly.param @U : !poly.tvar<@U> +// CHECK-NEXT: function.def @g(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !poly.tvar<@U>) -> !poly.tvar<@U> { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = function.call @NestedScopeInner::@f<[?, @U]>(%[[VAL_0]]) : (!poly.tvar<@U>) -> !poly.tvar<@U> +// CHECK-NEXT: function.return %[[VAL_1]] : !poly.tvar<@U> +// CHECK-NEXT: } +// CHECK-NEXT: } +// +// CHECK-LABEL: poly.template @NestedScopeCaller { +// CHECK-NEXT: poly.param @U : !poly.tvar<@U> +// CHECK-NEXT: function.def @call(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !poly.tvar<@U>) -> !poly.tvar<@U> { +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"NestedScopeOuter_i_\1A"::@g<[?]>(%[[VAL_2]]) : (!poly.tvar<@U>) -> !poly.tvar<@U> +// CHECK-NEXT: function.return %[[VAL_3]] : !poly.tvar<@U> +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + // Test transitive instantiation of functions. Also different instantiations of the same function. module attributes {llzk.lang} { poly.template @template_g1 { @@ -356,3 +405,51 @@ module attributes {llzk.lang} { // CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = function.call @T_89_i_compute() : () -> index // CHECK-NEXT: function.return %[[VAL_16]] : index // CHECK-NEXT: } +// ----- + +// Reusing a full instantiation across different call-site spellings must keep each rewritten +// callee valid from its caller scope. +module attributes {llzk.lang} { + module @lib attributes {llzk.lang} { + poly.template @T { + poly.param @N : index + + function.def @f(%inp: !array.type<@N x !felt.type>) -> !felt.type { + %c0 = arith.constant 0 : index + %r = array.read %inp[%c0] : !array.type<@N x !felt.type>, !felt.type + function.return %r : !felt.type + } + } + + function.def @local(%inp: !array.type<4 x !felt.type>) -> !felt.type { + %r = function.call @T::@f(%inp) : (!array.type<4 x !felt.type>) -> !felt.type + function.return %r : !felt.type + } + } + + module @sibling { + function.def @remote(%inp: !array.type<4 x !felt.type>) -> !felt.type { + %r = function.call @lib::@T::@f(%inp) : (!array.type<4 x !felt.type>) -> !felt.type + function.return %r : !felt.type + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: module @lib attributes {llzk.lang} { +// CHECK-NEXT: function.def @T_4_f(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_0]]{{\[}}%[[VAL_1]]] : <4 x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_2]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: function.def @local(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = function.call @T_4_f(%[[VAL_3]]) : (!array.type<4 x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_4]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: module @sibling { +// CHECK-NEXT: function.def @remote(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !array.type<4 x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @lib::@T_4_f(%[[VAL_5]]) : (!array.type<4 x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_6]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/instantiate_name_encoding.llzk b/test/Transforms/Flattening/instantiate_name_encoding.llzk index 0edd8af217..2a773ed3fd 100644 --- a/test/Transforms/Flattening/instantiate_name_encoding.llzk +++ b/test/Transforms/Flattening/instantiate_name_encoding.llzk @@ -40,6 +40,68 @@ module attributes {llzk.lang} { // ----- +module attributes {llzk.lang} { + // These two distinct type tuples render identically because BuildShortTypeString leaves symbol + // names unescaped. Full-instantiation reuse must use the exact bindings, not that rendered name. + struct.def @a { + function.def @compute() -> !struct.type<@a> { + %self = struct.new : !struct.type<@a> + function.return %self : !struct.type<@a> + } + function.def @constrain(%self: !struct.type<@a>) { function.return } + } + struct.def @c { + function.def @compute() -> !struct.type<@c> { + %self = struct.new : !struct.type<@c> + function.return %self : !struct.type<@c> + } + function.def @constrain(%self: !struct.type<@c>) { function.return } + } + struct.def @"a>_!s<@b" { + function.def @compute() -> !struct.type<@"a>_!s<@b"> { + %self = struct.new : !struct.type<@"a>_!s<@b"> + function.return %self : !struct.type<@"a>_!s<@b"> + } + function.def @constrain(%self: !struct.type<@"a>_!s<@b">) { function.return } + } + struct.def @"b>_!s<@c" { + function.def @compute() -> !struct.type<@"b>_!s<@c"> { + %self = struct.new : !struct.type<@"b>_!s<@c"> + function.return %self : !struct.type<@"b>_!s<@c"> + } + function.def @constrain(%self: !struct.type<@"b>_!s<@c">) { function.return } + } + + poly.template @T { + poly.param @A : !poly.tvar<@A> + poly.param @B : !poly.tvar<@B> + function.def @value() -> index { + %value = arith.constant 7 : index + function.return %value : index + } + } + + function.def @main() -> index { + %first = function.call @T::@value<[!struct.type<@"a>_!s<@b">, !struct.type<@c>]>() : () -> index + %second = function.call @T::@value<[!struct.type<@a>, !struct.type<@"b>_!s<@c">]>() : () -> index + function.return %second : index + } +} + +// CHECK-LABEL: function.def @"T_!s<@a>_!s<@b>_!s<@c>_value"() -> index { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 7 : index +// CHECK-NEXT: function.return %[[VAL_0]] : index +// CHECK-NEXT: } +// CHECK-LABEL: function.def @"T_!s<@a>_!s<@b>_!s<@c>_value_0"() -> index { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = arith.constant 7 : index +// CHECK-NEXT: function.return %[[VAL_1]] : index +// CHECK-NEXT: } +// CHECK-LABEL: function.def @main() -> index { +// CHECK: %[[VAL_2:[0-9a-zA-Z_\.]+]] = function.call @"T_!s<@a>_!s<@b>_!s<@c>_value"() : () -> index +// CHECK: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"T_!s<@a>_!s<@b>_!s<@c>_value_0"() : () -> index + +// ----- + module attributes {llzk.lang} { // A full instantiation renders concrete values in parameter order. poly.template @T { diff --git a/test/Transforms/Flattening/instantiate_structs_pass_1.llzk b/test/Transforms/Flattening/instantiate_structs_pass_1.llzk index 706b946790..e806585432 100644 --- a/test/Transforms/Flattening/instantiate_structs_pass_1.llzk +++ b/test/Transforms/Flattening/instantiate_structs_pass_1.llzk @@ -88,6 +88,35 @@ module attributes {llzk.lang} { // CHECK-NEXT: } // ----- +// Preserve discardable attributes when Step 1 instantiates a nondet result type. +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @instantiate_nondet() attributes {function.allow_witness} { + %witness = llzk.nondet : !struct.type<@TCell::@Cell<[0]>> {product_source = "direct"} + function.return + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK: function.def @instantiate_nondet() attributes {function.allow_witness} { +// CHECK-NEXT: %{{[0-9a-zA-Z_\.]+}} = llzk.nondet : !struct.type<@TCell_0_Cell> {product_source = "direct"} +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// ----- + module attributes {llzk.lang} { poly.template @TComponent01A { poly.param @A diff --git a/test/Transforms/Flattening/instantiate_wildcard.llzk b/test/Transforms/Flattening/instantiate_wildcard.llzk index 1351fdcc80..0427c956db 100644 --- a/test/Transforms/Flattening/instantiate_wildcard.llzk +++ b/test/Transforms/Flattening/instantiate_wildcard.llzk @@ -65,7 +65,7 @@ module attributes {llzk.lang = "circom", llzk.main = !struct.type<@CallDiffTypeT %felt_const_0_0 = felt.const 0 : <"bn128"> %1 = cast.toindex %felt_const_0_0 : !felt.type<"bn128"> %2 = poly.unifiable_cast %arg0 : (!poly.tvar<@T_arg0>) -> !array.type> - %3 = array.read %2[%0, %1] : >, !poly.tvar<@"$e"> + %3 = array.read %2[%0, %1] : >, !poly.tvar<@"$e"> {product_source = "read"} %4 = poly.unifiable_cast %3 : (!poly.tvar<@"$e">) -> !poly.tvar<@T_return> function.return %4 : !poly.tvar<@T_return> } @@ -109,7 +109,7 @@ module attributes {llzk.lang = "circom", llzk.main = !struct.type<@CallDiffTypeT // CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = cast.toindex %[[VAL_1]] : !felt.type<"bn128"> // CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = cast.toindex %[[VAL_1]] : !felt.type<"bn128"> // CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_0]] : (!array.type<10,5,5 x !felt.type<"bn128">>) -> !array.type> -// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.extract %[[VAL_4]]{{\[}}%[[VAL_2]], %[[VAL_3]]] : > +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.extract %[[VAL_4]]{{\[}}%[[VAL_2]], %[[VAL_3]]] : > {product_source = "read"} // CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_5]] : (!array.type<5 x !felt.type<"bn128">>) -> !array.type<5 x !felt.type<"bn128">> // CHECK-NEXT: function.return %[[VAL_6]] : !array.type<5 x !felt.type<"bn128">> // CHECK-NEXT: } diff --git a/test/Transforms/Flattening/member_read_array_initializer_call_consumer_guard.llzk b/test/Transforms/Flattening/member_read_array_initializer_call_consumer_guard.llzk new file mode 100644 index 0000000000..c37d62cf79 --- /dev/null +++ b/test/Transforms/Flattening/member_read_array_initializer_call_consumer_guard.llzk @@ -0,0 +1,48 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume1(%values: !array.type<1 x !struct.type<@TCell::@Cell<[1]>>>) { + function.return + } + + // Refining the member read to Cell<0> makes it a concrete array.new initializer. The array itself + // must stay generic because the cast feeding this call independently requires array>. + struct.def @Holder { + struct.member @value : !struct.type<@TCell::@Cell<[#id]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@value] = %item0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + %read = struct.readm %self[@value] : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + %local = array.new %read : !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-error@+1 {{'poly.unifiable_cast' op input type '!array.type<1 x !struct.type<@TCell::@Cell<[affine_map<(d0) -> (d0)>]>>>' and output type '!array.type<1 x !struct.type<@TCell_1_Cell>>' are not unifiable}} + %as1 = poly.unifiable_cast %local : (!array.type<1 x !struct.type<@TCell::@Cell<[#id]>>>) -> !array.type<1 x !struct.type<@TCell::@Cell<[1]>>> + function.call @consume1(%as1) : (!array.type<1 x !struct.type<@TCell::@Cell<[1]>>>) -> () + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/member_read_array_write_consumer_guard.llzk b/test/Transforms/Flattening/member_read_array_write_consumer_guard.llzk new file mode 100644 index 0000000000..410a0f9396 --- /dev/null +++ b/test/Transforms/Flattening/member_read_array_write_consumer_guard.llzk @@ -0,0 +1,48 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + // A generic member read may unify with the write-derived member refinement and with the + // independently concrete element type of an array.write destination. Keep the member generic + // unless that destination accepts the proposed refinement. + struct.def @Holder { + struct.member @value : !struct.type<@TCell::@Cell<[#id]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@value] = %item0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + %read = struct.readm %self[@value] : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + %index = arith.constant 0 : index + %dest = array.new : !array.type<1 x !struct.type<@TCell::@Cell<[1]>>> + // The generic read remains unchanged, so this independently concrete destination is + // diagnosed instead of silently retagging the read to Cell<0>. + // expected-error@+1 {{'array.write' op failed to verify that rvalue type matches with arr_ref element type}} + array.write %dest[%index] = %read : !array.type<1 x !struct.type<@TCell::@Cell<[1]>>>, !struct.type<@TCell::@Cell<[#id]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/member_shared_nondet_external_constraint.llzk b/test/Transforms/Flattening/member_shared_nondet_external_constraint.llzk new file mode 100644 index 0000000000..3cdd6fb252 --- /dev/null +++ b/test/Transforms/Flattening/member_shared_nondet_external_constraint.llzk @@ -0,0 +1,60 @@ +// RUN: llzk-opt -llzk-flatten %s | FileCheck %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @item : !struct.type<@TCell::@Cell<[#id]>> {column} + + function.def @product() -> !struct.type<@Holder> + attributes {function.allow_witness, function.allow_constraint} { + %self = struct.new : !struct.type<@Holder> + %shared = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> {product_source = "shared"} + function.call @TCell::@Cell::@constrain(%shared) : (!struct.type<@TCell::@Cell<[#id]>>) -> () {product_source = "shared"} + struct.writem %self[@item] = %shared : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + %concrete = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@item] = %concrete : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + function.return %self : !struct.type<@Holder> + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @item : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @product() -> !struct.type<@Holder> attributes {function.allow_constraint, function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> {product_source = "shared"} +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_3]]) : (!struct.type<@TCell_0_Cell>) -> () {product_source = "shared"} +// CHECK-NEXT: struct.writem %[[VAL_2]][@item] = %[[VAL_3]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_2]][@item] = %[[VAL_4]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/partial_function_expr.llzk b/test/Transforms/Flattening/partial_function_expr.llzk new file mode 100644 index 0000000000..24779e1d5d --- /dev/null +++ b/test/Transforms/Flattening/partial_function_expr.llzk @@ -0,0 +1,249 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten %s | FileCheck --enable-var-scope %s + +// A partial function instantiation must retain expressions that read the remaining parameters. +module attributes {llzk.lang} { + poly.template @Callee { + poly.param @A : index + poly.param @B : index + poly.expr @Length { + %a = poly.read_const @A : index + %b = poly.read_const @B : index + %length = arith.addi %a, %b : index + poly.yield %length : index + } + + function.def @f(%input: !array.type<@A,@B x !felt.type>) -> !felt.type { + %zero = arith.constant 0 : index + %length = poly.read_const @Length : index + %result = array.read %input[%zero, %zero] : !array.type<@A,@B x !felt.type>, !felt.type + function.return %result : !felt.type + } + } + + poly.template @Caller { + poly.param @N : index + + function.def @g(%input: !array.type<8,@N x !felt.type>) -> !felt.type { + %result = function.call @Callee::@f(%input) : (!array.type<8,@N x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: poly.template @"Callee_8_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: poly.expr @Length { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @B : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index +// CHECK-NEXT: poly.yield %[[VAL_2]] : index +// CHECK-NEXT: } +// CHECK-NEXT: function.def @f(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<8,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_3]]{{\[}}%[[VAL_4]], %[[VAL_4]]] : <8,@B x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_5]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } {poly.name_pattern = ["Callee_8_", ""]} +// CHECK-NEXT: poly.template @Caller { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: function.def @g(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<8,@N x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@f(%[[VAL_6]]) : (!array.type<8,@N x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_7]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// A partial function instantiation must retain flat calls to sibling functions in the source +// template by copying the referenced sibling into the partial template. +module attributes {llzk.lang} { + poly.template @Callee { + poly.param @A : index + poly.param @B : index + + function.def @leaf(%input: !array.type<@A,@B x !felt.type>) -> !felt.type { + %zero = arith.constant 0 : index + %result = array.read %input[%zero, %zero] : !array.type<@A,@B x !felt.type>, !felt.type + function.return %result : !felt.type + } + + function.def @helper(%input: !array.type<@A,@B x !felt.type>) -> !felt.type { + %result = function.call @Callee::@leaf(%input) : (!array.type<@A,@B x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + + function.def @f(%input: !array.type<@A,@B x !felt.type>) -> !felt.type { + %result = function.call @Callee::@helper(%input) : (!array.type<@A,@B x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + } + + poly.template @Caller { + poly.param @N : index + + function.def @g(%input: !array.type<8,@N x !felt.type>) -> !felt.type { + %result = function.call @Callee::@f(%input) : (!array.type<8,@N x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: poly.template @"Callee_8_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: function.def @f(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<8,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@helper(%[[VAL_0]]) : (!array.type<8,@B x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_1]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: function.def @helper(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<8,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@leaf(%[[VAL_2]]) : (!array.type<8,@B x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_3]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: function.def @leaf(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<8,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_4]]{{\[}}%[[VAL_5]], %[[VAL_5]]] : <8,@B x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_6]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } {poly.name_pattern = ["Callee_8_", ""]} +// CHECK-NEXT: poly.template @Callee { +// CHECK-NEXT: poly.param @A : index +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: function.def @leaf(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !array.type<@A,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_7]]{{\[}}%[[VAL_8]], %[[VAL_8]]] : <@A,@B x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_9]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: function.def @helper(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !array.type<@A,@B x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @Callee::@leaf(%[[VAL_10]]) : (!array.type<@A,@B x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_11]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @Caller { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: function.def @g(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !array.type<8,@N x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@f(%[[VAL_12]]) : (!array.type<8,@N x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_13]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// A partial function instantiation must retain expressions referenced only from types/attributes. +module attributes {llzk.lang} { + poly.template @Callee { + poly.param @A : index + poly.param @B : index + poly.expr @Length { + %a = poly.read_const @A : index + %b = poly.read_const @B : index + %length = arith.addi %a, %b : index + poly.yield %length : index + } + + function.def @f(%input: !array.type<@A,@B x !felt.type>) -> !felt.type attributes {test.signature_type = !array.type<@Length x !felt.type>} { + %zero = arith.constant 0 : index + %result = array.read %input[%zero, %zero] : !array.type<@A,@B x !felt.type>, !felt.type + function.return %result : !felt.type + } + } + + poly.template @Caller { + poly.param @N : index + + function.def @g(%input: !array.type<8,@N x !felt.type>) -> !felt.type { + %result = function.call @Callee::@f(%input) : (!array.type<8,@N x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: poly.template @"Callee_8_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: poly.expr @Length { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @B : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index +// CHECK-NEXT: poly.yield %[[VAL_2]] : index +// CHECK-NEXT: } +// CHECK-NEXT: function.def @f(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<8,@B x !felt.type>) -> !felt.type attributes {test.signature_type = !array.type<@Length x !felt.type>} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_3]]{{\[}}%[[VAL_4]], %[[VAL_4]]] : <8,@B x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_5]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } {poly.name_pattern = ["Callee_8_", ""]} +// CHECK-NEXT: poly.template @Caller { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: function.def @g(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !array.type<8,@N x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@f(%[[VAL_6]]) : (!array.type<8,@N x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_7]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// A partial function instantiation must retain `poly.expr` referenced only by its signature. +module attributes {llzk.lang} { + poly.template @Callee { + poly.param @A : index + poly.param @B : index + poly.expr @Length { + %a = poly.read_const @A : index + %b = poly.read_const @B : index + %length = arith.addi %a, %b : index + poly.yield %length : index + } + + function.def @f(%input: !array.type<@Length x !felt.type>) -> !felt.type { + %zero = arith.constant 0 : index + %result = array.read %input[%zero] : !array.type<@Length x !felt.type>, !felt.type + function.return %result : !felt.type + } + } + + poly.template @Caller { + poly.param @N : index + poly.expr @Length { + %eight = arith.constant 8 : index + %n = poly.read_const @N : index + %length = arith.addi %eight, %n : index + poly.yield %length : index + } + + function.def @g(%input: !array.type<@Length x !felt.type>) -> !felt.type { + %result = function.call @Callee::@f<[8, @N]>(%input) : + (!array.type<@Length x !felt.type>) -> !felt.type + function.return %result : !felt.type + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: poly.template @"Callee_8_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: poly.expr @Length { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = poly.read_const @B : index +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_1]], %[[VAL_0]] : index +// CHECK-NEXT: poly.yield %[[VAL_2]] : index +// CHECK-NEXT: } +// CHECK-NEXT: function.def @f(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<@Length x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_3]]{{\[}}%[[VAL_4]]] : <@Length x !felt.type>, !felt.type +// CHECK-NEXT: function.return %[[VAL_5]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } {poly.name_pattern = ["Callee_8_", ""]} +// CHECK-NEXT: poly.template @Caller { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: poly.expr @Length { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = arith.constant 8 : index +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = poly.read_const @N : index +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = arith.addi %[[VAL_7]], %[[VAL_6]] : index +// CHECK-NEXT: poly.yield %[[VAL_8]] : index +// CHECK-NEXT: } +// CHECK-NEXT: function.def @g(%[[VAL_9:[0-9a-zA-Z_\.]+]]: !array.type<@Length x !felt.type>) -> !felt.type { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @"Callee_8_\1A"::@f<[@N]>(%[[VAL_9]]) : (!array.type<@Length x !felt.type>) -> !felt.type +// CHECK-NEXT: function.return %[[VAL_10]] : !felt.type +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/refinement_consumer_type_guards.llzk b/test/Transforms/Flattening/refinement_consumer_type_guards.llzk new file mode 100644 index 0000000000..0ba554aa1d --- /dev/null +++ b/test/Transforms/Flattening/refinement_consumer_type_guards.llzk @@ -0,0 +1,51 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck %s + +module attributes {llzk.lang} { + function.def @consume3x3(%value: !array.type<3,3 x !felt.type>) { + function.return + } + + struct.def @Holder { + struct.member @items : !array.type {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %source = array.new : !array.type<2,2 x !felt.type> + %first = poly.unifiable_cast %source : (!array.type<2,2 x !felt.type>) -> !array.type<2,? x !felt.type> + %second = poly.unifiable_cast %source : (!array.type<2,2 x !felt.type>) -> !array.type + // The complementary writes refine @items to 2 x 2. + struct.writem %self[@items] = %first : !struct.type<@Holder>, !array.type<2,? x !felt.type> + struct.writem %self[@items] = %second : !struct.type<@Holder>, !array.type + // This generic read instead feeds a call whose declared parameter requires 3 x 3. + %read = struct.readm %self[@items] : !struct.type<@Holder>, !array.type + function.call @consume3x3(%read) : (!array.type) -> () + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: function.def @consume3x3(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<3,3 x !felt.type>) { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items : !array.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = array.new : <2,2 x !felt.type> +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_2]] : (!array.type<2,2 x !felt.type>) -> !array.type +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_2]] : (!array.type<2,2 x !felt.type>) -> !array.type +// CHECK-NEXT: struct.writem %[[VAL_1]][@items] = %[[VAL_3]] : <@Holder>, !array.type +// CHECK-NEXT: struct.writem %[[VAL_1]][@items] = %[[VAL_4]] : <@Holder>, !array.type +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_1]][@items] : <@Holder>, !array.type +// CHECK-NEXT: function.call @consume3x3(%[[VAL_5]]) : (!array.type) -> () +// CHECK-NEXT: function.return %[[VAL_1]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/refinement_consumer_type_guards_fail.llzk b/test/Transforms/Flattening/refinement_consumer_type_guards_fail.llzk new file mode 100644 index 0000000000..acc09323ac --- /dev/null +++ b/test/Transforms/Flattening/refinement_consumer_type_guards_fail.llzk @@ -0,0 +1,441 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // The generic read is valid before scalarization, but its concrete index-zero value cannot + // replace this use, which requires Cell<1>. + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a member write}} + %read0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{member write requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@wrong] = %read0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} + +// ----- + +// A generic array read can feed a concrete unifiable cast before scalarization. Replacing the +// read with Cell<0> must be rejected because the cast independently requires Cell<1>. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a unifiable cast}} + %read0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{unifiable cast result requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + %cast = poly.unifiable_cast %read0 : (!struct.type<@TCell::@Cell<[#id]>>) -> !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@wrong] = %cast : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %source : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} + +// ----- + +// Regression: refreshed candidates must validate read consumers before any candidate rewrites. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + function.def @refreshedCandidateConsumer() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %first = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %first[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %first[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %generic0 = array.read %first[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + + // Refreshing this candidate makes index zero Cell<0>; its declared generic read still + // verifies, but the concrete consumer below requires Cell<1>. + %second = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %second[%c0] = %generic0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %second[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a function call}} + %read0 = array.read %second[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{call argument requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + function.call @consume1(%read0) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + function.return + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @arrayInitializerConsumer() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // This generic read is a valid array initializer before scalarization, but the concrete + // index-zero replacement cannot exactly match the destination's generic element type. + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with an array initializer}} + %read0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{array initializer requires exactly '!struct.type<@TCell::@Cell<[affine_map<(d0) -> (d0)>]>>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + %dest = array.new %read0 : !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> + function.return + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @arrayWriteConsumer() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with an array write}} + %read0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + %dest = array.new : !array.type<1 x !struct.type<@TCell::@Cell<[1]>>> + // expected-note@+1 {{array write requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + array.write %dest[%c0] = %read0 : !array.type<1 x !struct.type<@TCell::@Cell<[1]>>>, !struct.type<@TCell::@Cell<[#id]>> + function.return + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @returnConsumer() -> !struct.type<@TCell::@Cell<[1]>> attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a function return}} + %read0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{function return requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + function.return %read0 : !struct.type<@TCell::@Cell<[#id]>> + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + struct.def @Cell { + struct.member @value : !felt.type {column} + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // The generic read is valid before scalarization, but its concrete index-zero value cannot + // replace this argument, which requires Cell<1>. + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a function call}} + %read0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{call argument requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + function.call @consume1(%read0) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + %c0 = arith.constant 0 : index + %items = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-error@+1 {{cannot scalarize heterogeneous array read because its index-specific value type is incompatible with a function call}} + %item0 = array.read %items[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // expected-note@+1 {{call argument requires '!struct.type<@TCell_1_Cell>', but this read is replaced with '!struct.type<@TCell_0_Cell>'}} + function.call @consume1(%item0) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + // expected-error@+1 {{'array.write' op failed to verify that rvalue type matches with arr_ref element type}} + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %wrong = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@wrong] = %wrong : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} + +// ----- + +// A write proposes Cell<0> for @item, while this public member read is returned through a +// function declared as Cell<1>. The write-derived refinement must be rejected before the read is +// retagged to Cell<0>. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @item : !struct.type<@TCell::@Cell<[#id]>> {column, llzk.pub} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@item] = %item : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } + + function.def @readItem(%self: !struct.type<@Holder>) -> !struct.type<@TCell::@Cell<[1]>> { + %item = struct.readm %self[@item] : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{type of return operand 0 ('!struct.type<@TCell_0_Cell>') doesn't match function result type ('!struct.type<@TCell_1_Cell>') in function @readItem}} + function.return %item : !struct.type<@TCell::@Cell<[#id]>> + } +} diff --git a/test/Transforms/Flattening/scalarize_array_member_read_offset.llzk b/test/Transforms/Flattening/scalarize_array_member_read_offset.llzk new file mode 100644 index 0000000000..1fefdb7a75 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_array_member_read_offset.llzk @@ -0,0 +1,95 @@ +// RUN: llzk-opt -llzk-flatten %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column, product_source = "member"} + + function.def @compute(%idx: index) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %items = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %items[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> {product_source = "item_0", write_source = "element"} + array.write %items[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> {product_source = "item_1", write_source = "element"} + struct.writem %self[@items] = %items : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {product_source = "compute", member_source = "whole-array"} + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %idx: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %prev_items = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {product_source = "constrain", tableOffset = -1 : index} + %prev_item = array.read %prev_items[%c1] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %affine_items = struct.readm %self[@items] {()[%idx]} : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {product_source = "constrain", tableOffset = affine_map<()[s0] -> (s0 - 1)>} + %affine_item = array.read %affine_items[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %prev_value = struct.readm %prev_item[@value] : !struct.type<@TCell::@Cell<[1]>>, !felt.type + %affine_value = struct.readm %affine_item[@value] : !struct.type<@TCell::@Cell<[0]>>, !felt.type + constrain.eq %prev_value, %affine_value : !felt.type, !felt.type + function.return + } + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<()[s0] -> (s0 - 1)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column, product_source = "member"} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column, product_source = "member"} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: index) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> {member_source = "whole-array", product_source = "item_0", write_source = "element"} +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> {member_source = "whole-array", product_source = "item_1", write_source = "element"} +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_9:[0-9a-zA-Z_\.]+]]: index) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_8]][@items_1] : <@Holder>, !struct.type<@TCell_1_Cell> {product_source = "constrain", tableOffset = -1 : index} +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_8]][@items_0] {(){{\[}}%[[VAL_9]]]} : <@Holder>, !struct.type<@TCell_0_Cell> {product_source = "constrain", tableOffset = #[[$ATTR_0]]} +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_10]][@value] : <@TCell_1_Cell>, !felt.type +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_11]][@value] : <@TCell_0_Cell>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_12]], %[[VAL_13]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_dead_cross_block_writes.llzk b/test/Transforms/Flattening/scalarize_dead_cross_block_writes.llzk new file mode 100644 index 0000000000..3224cebcab --- /dev/null +++ b/test/Transforms/Flattening/scalarize_dead_cross_block_writes.llzk @@ -0,0 +1,78 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + scf.if %cond { + %c0 = arith.constant 0 : index + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + } else { + %c1 = arith.constant 1 : index + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_9:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_dependent_arrays.llzk b/test/Transforms/Flattening/scalarize_dependent_arrays.llzk new file mode 100644 index 0000000000..f4219c8055 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_dependent_arrays.llzk @@ -0,0 +1,84 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %from_source0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %from_source1 = array.read %source[%c1] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %copy = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %copy[%c0] = %from_source0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %copy[%c1] = %from_source1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %copy : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_5]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_external_constraint_generic_signature.llzk b/test/Transforms/Flattening/scalarize_external_constraint_generic_signature.llzk new file mode 100644 index 0000000000..d30423abea --- /dev/null +++ b/test/Transforms/Flattening/scalarize_external_constraint_generic_signature.llzk @@ -0,0 +1,67 @@ +// RUN: llzk-opt -llzk-flatten %s | FileCheck %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute(%values: !array.type<@N x !felt.type>) -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self: !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>, %values: !array.type<@N x !felt.type>) { + function.return + } + } + } + + function.def @generic_constraint_signature(%values: !array.type<#id x !felt.type>) + attributes {function.allow_witness, function.allow_constraint} { + %c0 = arith.constant 0 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + function.call @TCell::@Cell::@constrain(%generic, %values) : (!struct.type<@TCell::@Cell<[#id]>>, !array.type<#id x !felt.type>) -> () + + %other = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %local = array.new %generic, %other : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.return + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute(%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<0 x !felt.type>) -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_1:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_1]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_2:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>, %[[VAL_3:[0-9a-zA-Z_\.]+]]: !array.type<0 x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<@N x !felt.type>) -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>, %[[VAL_7:[0-9a-zA-Z_\.]+]]: !array.type<@N x !felt.type>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: function.def @generic_constraint_signature(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !array.type<#[[$ATTR_0]] x !felt.type>) attributes {function.allow_constraint, function.allow_witness} { +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_10]], %[[VAL_8]]) : (!struct.type<@TCell_0_Cell>, !array.type<#[[$ATTR_0]] x !felt.type>) -> () +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_external_constraint_signature_fail.llzk b/test/Transforms/Flattening/scalarize_external_constraint_signature_fail.llzk new file mode 100644 index 0000000000..8130dfe071 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_external_constraint_signature_fail.llzk @@ -0,0 +1,38 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute(%values: !array.type<@N x !felt.type>) -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self: !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>, %values: !array.type<@N x !felt.type>) { + function.return + } + } + } + + function.def @constraint_signature_conflict() attributes {function.allow_witness, function.allow_constraint} { + %c0 = arith.constant 0 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %values = array.new : !array.type<1 x !felt.type> + + // Both operands independently unify with the generic constraint signature. Specializing + // %generic to Cell<0> must not retarget this call, however, because %values remains array<1>. + // expected-note@+1 {{retargeted constraint expects argument types '!struct.type<@TCell_0_Cell>', '!array.type<0 x !felt.type>'}} + function.call @TCell::@Cell::@constrain(%generic, %values) : (!struct.type<@TCell::@Cell<[#id]>>, !array.type<1 x !felt.type>) -> () + + %other = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{cannot scalarize array because specializing a generic nondeterministic initializer would make an external constraint call incompatible with its retargeted callee}} + %local = array.new %generic, %other : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.return + } +} diff --git a/test/Transforms/Flattening/scalarize_inline_initializer.llzk b/test/Transforms/Flattening/scalarize_inline_initializer.llzk new file mode 100644 index 0000000000..d81bc7bdcf --- /dev/null +++ b/test/Transforms/Flattening/scalarize_inline_initializer.llzk @@ -0,0 +1,412 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %generic0 = poly.unifiable_cast %item0 : (!struct.type<@TCell::@Cell<[0]>>) -> !struct.type<@TCell::@Cell<[#id]>> + %generic1 = poly.unifiable_cast %item1 : (!struct.type<@TCell::@Cell<[1]>>) -> !struct.type<@TCell::@Cell<[#id]>> + %local = array.new %generic0, %generic1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + scf.if %cond { + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_6]] : (!struct.type<@TCell_0_Cell>) -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_7]] : (!struct.type<@TCell_1_Cell>) -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %source = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %from_source0 = array.read %source[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + %from_source1 = array.read %source[%c1] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + %copy = array.new %from_source0, %from_source1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@items] = %copy : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_5]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + poly.template @Any { + poly.param @T : !poly.tvar<@T> + function.def @f(%a: !poly.tvar<@T>) -> !poly.tvar<@T> { + function.return %a : !poly.tvar<@T> + } + } + + function.def @use_inline_read_only_init() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + // Build an inline initializer whose stored SSA values stay generic while each element type + // specializes per static index during scalarization. + %generic0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %generic1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %local = array.new %generic0, %generic1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + // Regression: if scalarization replaces these reads with %generic0 directly, the writes below are + // ill-typed. If it creates one specialized nondet per read, the original equality between repeated + // reads of %local[%c0] is lost. Both reads must use one concrete index-specialized replacement value. + %picked0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %picked1 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %sink = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[0]>>> + array.write %sink[%c0] = %picked0 : !array.type<2 x !struct.type<@TCell::@Cell<[0]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %sink[%c1] = %picked1 : !array.type<2 x !struct.type<@TCell::@Cell<[0]>>>, !struct.type<@TCell::@Cell<[0]>> + + function.call @Any::@f(%sink) : (!array.type<2 x !struct.type<@TCell::@Cell<[0]>>>) -> !array.type<2 x !struct.type<@TCell::@Cell<[0]>>> + + function.return + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: function.def @"Any_!a:2>_f"(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !array.type<2 x !struct.type<@TCell_0_Cell>>) -> !array.type<2 x !struct.type<@TCell_0_Cell>> { +// CHECK-NEXT: function.return %[[VAL_4]] : !array.type<2 x !struct.type<@TCell_0_Cell>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @use_inline_read_only_init() attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = array.new : <2 x !struct.type<@TCell_0_Cell>> +// CHECK-NEXT: array.write %[[VAL_10]]{{\[}}%[[VAL_5]]] = %[[VAL_9]] : <2 x !struct.type<@TCell_0_Cell>>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: array.write %[[VAL_10]]{{\[}}%[[VAL_6]]] = %[[VAL_9]] : <2 x !struct.type<@TCell_0_Cell>>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @"Any_!a:2>_f"(%[[VAL_10]]) : (!array.type<2 x !struct.type<@TCell_0_Cell>>) -> !array.type<2 x !struct.type<@TCell_0_Cell>> +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @shared_initializer_sibling_reads(%cond: i1) attributes {function.allow_witness, function.allow_constraint} { + %c0 = arith.constant 0 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> {product_source = "generic"} + %local = array.new %generic, %generic : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + scf.if %cond { + %then = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.call @TCell::@Cell::@constrain(%then) : (!struct.type<@TCell::@Cell<[0]>>) -> () + } else { + %else = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.call @TCell::@Cell::@constrain(%else) : (!struct.type<@TCell::@Cell<[0]>>) -> () + } + function.return + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: function.def @shared_initializer_sibling_reads(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint, function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> {product_source = "generic"} +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> {product_source = "generic"} +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_6]]) : (!struct.type<@TCell_0_Cell>) -> () +// CHECK-NEXT: } else { +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_6]]) : (!struct.type<@TCell_0_Cell>) -> () +// CHECK-NEXT: } +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @shared_initializer_external_use() + attributes {function.allow_witness, function.allow_constraint} { + %c0 = arith.constant 0 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // This accepted external observation precedes the candidate array. Its replacement must be + // materialized where it dominates this call as well as the later scalarized read. + function.call @TCell::@Cell::@constrain(%generic) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + %other = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + + // One SSA witness is observed both as the generic initializer and through a concrete array + // read. Scalarization must retain that sharing after specializing the witness. + %local = array.new %generic, %other : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %concrete = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.call @TCell::@Cell::@constrain(%concrete) : (!struct.type<@TCell::@Cell<[0]>>) -> () + function.return + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: function.def @shared_initializer_external_use() attributes {function.allow_constraint, function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_5]]) : (!struct.type<@TCell_0_Cell>) -> () +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: function.call @TCell_0_Cell::@constrain(%[[VAL_5]]) : (!struct.type<@TCell_0_Cell>) -> () +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_inline_initializer_fail.llzk b/test/Transforms/Flattening/scalarize_inline_initializer_fail.llzk new file mode 100644 index 0000000000..b8fbea3f0f --- /dev/null +++ b/test/Transforms/Flattening/scalarize_inline_initializer_fail.llzk @@ -0,0 +1,45 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume0(%value: !struct.type<@TCell::@Cell<[0]>>) { + function.return + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + function.def @shared_initializer() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + + // expected-error@+1 {{cannot scalarize array because a shared nondeterministic initializer requires incompatible specialized types}} + %local = array.new %generic, %generic : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %first = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + // expected-note@-1 {{array index [0 : index] is read with specialized type '!struct.type<@TCell_0_Cell>'}} + // expected-note@+1 {{array index [1 : index] is read with specialized type '!struct.type<@TCell_1_Cell>'}} + %second = array.read %local[%c1] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + function.call @consume0(%first) : (!struct.type<@TCell::@Cell<[0]>>) -> () + function.call @consume1(%second) : (!struct.type<@TCell::@Cell<[1]>>) -> () + function.return + } +} diff --git a/test/Transforms/Flattening/scalarize_member_basic.llzk b/test/Transforms/Flattening/scalarize_member_basic.llzk new file mode 100644 index 0000000000..f5e3ce73de --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_basic.llzk @@ -0,0 +1,750 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @first : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %observed = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@first] = %observed : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @first : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@first] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %source = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %source0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %source1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %source2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %source0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %source1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c2] = %source2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %generic_first = array.read %source[%c0] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + + %local = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %generic_first : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %local[%c1] = %item0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c2] = %item1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_2 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_5]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_2] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %partial = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %partial[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@items] = %partial : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + %c0 = arith.constant 0 : index + %items = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item = array.read %items[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %value = struct.readm %item[@value] : !struct.type<@TCell::@Cell<[0]>>, !felt.type + constrain.eq %value, %value : !felt.type, !felt.type + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_10]][@items_0] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_12]][@value] : <@TCell_0_Cell>, !felt.type +// CHECK-NEXT: constrain.eq %[[VAL_13]], %[[VAL_13]] : !felt.type, !felt.type +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @first : !struct.type<@TCell::@Cell<[0]>> {column} + struct.member @second : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + %initial = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %initial0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %initial1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %initial[%c0] = %initial0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %initial[%c1] = %initial1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %initial : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %snapshot = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {product_source = "whole", whole_source = "whole"} + + %replacement = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %replacement0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %replacement1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %replacement[%c0] = %replacement0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %replacement[%c1] = %replacement1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %replacement : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %old0 = array.read %snapshot[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> {product_source = "read"} + %old0_again = array.read %snapshot[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> {product_source = "read-again"} + struct.writem %self[@first] = %old0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@second] = %old0_again : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @first : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @second : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_5]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_4]][@items_0] : <@Holder>, !struct.type<@TCell_0_Cell> {product_source = "read-again", whole_source = "whole"} +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_4]][@items_0] : <@Holder>, !struct.type<@TCell_0_Cell> {product_source = "read", whole_source = "whole"} +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@first] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@second] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_11:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TPair { + poly.param @A : index + poly.param @B : index + + struct.def @Pair { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TPair::@Pair<[@A, @B]>> { + %self = struct.new : !struct.type<@TPair::@Pair<[@A, @B]>> + function.return %self : !struct.type<@TPair::@Pair<[@A, @B]>> + } + + function.def @constrain(%self: !struct.type<@TPair::@Pair<[@A, @B]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> {column} + struct.member @first : !struct.type<@TPair::@Pair<[0, 0]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + scf.if %cond { + %local = array.new : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + %item0 = function.call @TPair::@Pair::@compute() : () -> !struct.type<@TPair::@Pair<[0, 0]>> + %partial0 = poly.unifiable_cast %item0 : (!struct.type<@TPair::@Pair<[0, 0]>>) -> !struct.type<@TPair::@Pair<[#id, 0]>> + %item1 = function.call @TPair::@Pair::@compute() : () -> !struct.type<@TPair::@Pair<[1, 1]>> + array.write %local[%c0] = %partial0 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[#id, 0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[1, 1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + } else { + %local = array.new : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + %item0 = function.call @TPair::@Pair::@compute() : () -> !struct.type<@TPair::@Pair<[0, 0]>> + %partial0 = poly.unifiable_cast %item0 : (!struct.type<@TPair::@Pair<[0, 0]>>) -> !struct.type<@TPair::@Pair<[0, #id]>> + %item1 = function.call @TPair::@Pair::@compute() : () -> !struct.type<@TPair::@Pair<[1, 1]>> + array.write %local[%c0] = %partial0 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[0, #id]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[1, 1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + } + + %items = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + %first = array.read %items[%c0] : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[0, 0]>> + struct.writem %self[@first] = %first : !struct.type<@Holder>, !struct.type<@TPair::@Pair<[0, 0]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TPair_0_0_Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TPair_0_0_Pair> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TPair_0_0_Pair> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TPair_0_0_Pair>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TPair_1_1_Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TPair_1_1_Pair> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TPair_1_1_Pair> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TPair_1_1_Pair> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TPair_1_1_Pair>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @"TPair_\1A_0" { +// CHECK-NEXT: poly.param @A : index +// CHECK-NEXT: struct.def @Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@"TPair_\1A_0"::@Pair<[@A]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@"TPair_\1A_0"::@Pair<[@A]>> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@"TPair_\1A_0"::@Pair<[@A]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@"TPair_\1A_0"::@Pair<[@A]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @"TPair_0_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: struct.def @Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@"TPair_0_\1A"::@Pair<[@B]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@"TPair_0_\1A"::@Pair<[@B]>> +// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@"TPair_0_\1A"::@Pair<[@B]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@"TPair_0_\1A"::@Pair<[@B]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TPair_0_0_Pair> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TPair_1_1_Pair> {column} +// CHECK-NEXT: struct.member @first : !struct.type<@TPair_0_0_Pair> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_8:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_8]] { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TPair_0_0_Pair::@compute() : () -> !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_10]] : (!struct.type<@TPair_0_0_Pair>) -> !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = function.call @TPair_1_1_Pair::@compute() : () -> !struct.type<@TPair_1_1_Pair> +// CHECK-NEXT: struct.writem %[[VAL_9]][@items_0] = %[[VAL_11]] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: struct.writem %[[VAL_9]][@items_1] = %[[VAL_12]] : <@Holder>, !struct.type<@TPair_1_1_Pair> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = function.call @TPair_0_0_Pair::@compute() : () -> !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_13]] : (!struct.type<@TPair_0_0_Pair>) -> !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: %[[VAL_15:[0-9a-zA-Z_\.]+]] = function.call @TPair_1_1_Pair::@compute() : () -> !struct.type<@TPair_1_1_Pair> +// CHECK-NEXT: struct.writem %[[VAL_9]][@items_0] = %[[VAL_14]] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: struct.writem %[[VAL_9]][@items_1] = %[[VAL_15]] : <@Holder>, !struct.type<@TPair_1_1_Pair> +// CHECK-NEXT: } +// CHECK-NEXT: %[[VAL_16:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_9]][@items_0] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: struct.writem %[[VAL_9]][@first] = %[[VAL_16]] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: function.return %[[VAL_9]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_17:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_18:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @first : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + + scf.if %cond { + %generic0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %first = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %first1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %first2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %first[%c0] = %generic0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %first[%c1] = %first1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %first[%c2] = %first2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@items] = %first : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %second = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %second0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %second1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %second2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %second[%c0] = %second0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %second[%c1] = %second1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %second[%c2] = %second2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@items] = %second : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } + + %snapshot = struct.readm %self[@items] : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %old0 = array.read %snapshot[%c0] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@first] = %old0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_2_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_2_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_2_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_2_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @items_2 : !struct.type<@TCell_2_Cell> {column} +// CHECK-NEXT: struct.member @first : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_6:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_6]] { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TCell_2_Cell::@compute() : () -> !struct.type<@TCell_2_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_2] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = function.call @TCell_2_Cell::@compute() : () -> !struct.type<@TCell_2_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_12]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_2] = %[[VAL_13]] : <@Holder>, !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_7]][@items_0] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@first] = %[[VAL_14]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_7]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_15:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_16:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_member_candidate_conflicts.llzk b/test/Transforms/Flattening/scalarize_member_candidate_conflicts.llzk new file mode 100644 index 0000000000..05470bd910 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_candidate_conflicts.llzk @@ -0,0 +1,404 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + // expected-error@+1 {{cannot split heterogeneous array member because candidate writes require incompatible scalar member types}} + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{candidate writes index [0 : index] with type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %swapped0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %swapped1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %fallback[%c0] = %swapped0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %fallback[%c1] = %swapped1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + // expected-note@+1 {{conflicting candidate writes the same index with type '!struct.type<@TCell_1_Cell>'}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +// Regression: candidate whole-array writes must agree on their exact index set because +// the size of the array member must be fixed, not dependent on input to the circuit. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + // expected-error@+1 {{cannot split heterogeneous array member because candidate whole-array writes use different index set}} + struct.member @items : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %short = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %short0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %short1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %short[%c0] = %short0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %short[%c1] = %short1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{candidate establishes 2 split member indices}} + struct.writem %self[@items] = %short : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %long = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %long0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %long1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %long2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %long[%c0] = %long0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %long[%c1] = %long1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %long[%c2] = %long2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + // expected-note@+1 {{candidate whole-array write has 3 array indices, including extra index [2 : index]}} + struct.writem %self[@items] = %long : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +// Regression: candidate whole-array writes must agree on their exact index set because +// the size of the array member must be fixed, not dependent on input to the circuit. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + // expected-error@+1 {{cannot split heterogeneous array member because candidate whole-array writes use different index set}} + struct.member @items : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> {column, llzk.pub} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %local = array.new{(%c2)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{candidate establishes 2 split member indices}} + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %fallback = array.new{(%c3)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + %fallback0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %fallback1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %fallback2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %fallback[%c0] = %fallback0 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %fallback[%c1] = %fallback1 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %fallback[%c2] = %fallback2 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + // expected-note@+1 {{candidate whole-array write has 3 array indices, including extra index [2 : index]}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @source : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + // expected-error@+1 {{cannot split heterogeneous array member because candidate writes require incompatible scalar member types}} + struct.member @items : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + + %source = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %source0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %source1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %source2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %source[%c0] = %source0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c1] = %source1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %source[%c2] = %source2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@source] = %source : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + + %read0 = array.read %source[%c0] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + %from_source = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %from_source1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %from_source2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %from_source[%c0] = %read0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %from_source[%c1] = %from_source1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %from_source[%c2] = %from_source2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{candidate writes index [0 : index] with type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@items] = %from_source : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + + scf.if %cond { + %concrete = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %concrete0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %concrete1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %concrete2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %concrete[%c0] = %concrete0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %concrete[%c1] = %concrete1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %concrete[%c2] = %concrete2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{conflicting candidate writes the same index with type '!struct.type<@TCell_1_Cell>'}} + struct.writem %self[@items] = %concrete : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } + + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + poly.template @THolder { + poly.param @Shared : !poly.tvar<@Shared> + + struct.def @Holder { + struct.member @items : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute( + %cond: i1, + // expected-note@+1 {{array index [1 : index] stores value type '!poly.tvar<@Shared>'}} + %shared: !poly.tvar<@Shared> + ) -> !struct.type<@THolder::@Holder<[@Shared]>> { + %self = struct.new : !struct.type<@THolder::@Holder<[@Shared]>> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %local = array.new : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %item3 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[3]>> + array.write %local[%c0] = %item0 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c2] = %item2 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + array.write %local[%c3] = %item3 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[3]>> + struct.writem %self[@items] = %local : !struct.type<@THolder::@Holder<[@Shared]>>, !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + // expected-error@+1 {{cannot split heterogeneous array member because a scalarization candidate uses a non-nondeterministic SSA value for a refined scalar member type}} + %fallback = array.new : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %item3 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[3]>> + array.write %fallback[%c0] = %item0 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %fallback[%c1] = %shared : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !poly.tvar<@Shared> + array.write %fallback[%c2] = %item2 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + array.write %fallback[%c3] = %item3 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[3]>> + // expected-note@+1 {{same index is written with scalar member type '!struct.type<@TCell_1_Cell>'}} + struct.writem %self[@items] = %fallback : !struct.type<@THolder::@Holder<[@Shared]>>, !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@THolder::@Holder<[@Shared]>> + } + + function.def @constrain(%self: !struct.type<@THolder::@Holder<[@Shared]>>, %cond: i1, %shared: !poly.tvar<@Shared>) { + function.return + } + } + } +} + +// ----- + +// Regression: all scalarization candidates must be checked against the combined +// split layout. The first candidate establishes concrete types at every index; +// the second reuses one generic value at two indices whose concrete types differ. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + poly.template @THolder { + poly.param @Shared : !poly.tvar<@Shared> + + struct.def @Holder { + struct.member @items : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1, %sharedInput: !poly.tvar<@Shared>) -> !struct.type<@THolder::@Holder<[@Shared]>> { + %self = struct.new : !struct.type<@THolder::@Holder<[@Shared]>> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %local = array.new : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %item3 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[3]>> + array.write %local[%c0] = %item0 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c2] = %item2 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + array.write %local[%c3] = %item3 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[3]>> + struct.writem %self[@items] = %local : !struct.type<@THolder::@Holder<[@Shared]>>, !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %shared = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{cannot split heterogeneous array member because a scalarization candidate reuses one SSA value for incompatible scalar member types}} + %fallback = array.new : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %item3 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[3]>> + array.write %fallback[%c0] = %shared : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %fallback[%c1] = %shared : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %fallback[%c2] = %item2 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + array.write %fallback[%c3] = %item3 : !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[3]>> + // expected-note@+2 {{value is used for scalar member type}} + // expected-note@+1 {{same value is also used for scalar member type}} + struct.writem %self[@items] = %fallback : !struct.type<@THolder::@Holder<[@Shared]>>, !array.type<4 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@THolder::@Holder<[@Shared]>> + } + + function.def @constrain(%self: !struct.type<@THolder::@Holder<[@Shared]>>, %cond: i1, %sharedInput: !poly.tvar<@Shared>) { + function.return + } + } + } +} diff --git a/test/Transforms/Flattening/scalarize_member_expandable.llzk b/test/Transforms/Flattening/scalarize_member_expandable.llzk new file mode 100644 index 0000000000..bd65ccc358 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_expandable.llzk @@ -0,0 +1,415 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @observed : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %generic0 = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@observed] = %generic0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @observed : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@observed] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %left0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %left1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %right = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %right0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %right1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @left_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @left_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @right_2 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @right_3 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_13:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @observed : !struct.type<@TCell::@Cell<[#id]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %local[%c0] = %generic : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + // Regression: the read is rewritten before the later whole-array member write requests + // the concrete Cell<0> refinement. Both operations originally observe %generic, so their + // scalar replacements must continue to share one witness. + %read0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@observed] = %read0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @observed : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_6:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_6]] { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@observed] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_12]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_7]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_13:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_14:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TPair { + poly.param @A : index + poly.param @B : index + + struct.def @Pair { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TPair::@Pair<[@A, @B]>> { + %self = struct.new : !struct.type<@TPair::@Pair<[@A, @B]>> + function.return %self : !struct.type<@TPair::@Pair<[@A, @B]>> + } + + function.def @constrain(%self: !struct.type<@TPair::@Pair<[@A, @B]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !struct.type<@TPair::@Pair<[#id, 0]>> {column} + struct.member @right : !struct.type<@TPair::@Pair<[0, #id]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %shared = llzk.nondet : !struct.type<@TPair::@Pair<[#id, #id]>> + struct.writem %self[@left] = %shared : !struct.type<@Holder>, !struct.type<@TPair::@Pair<[#id, #id]>> + struct.writem %self[@right] = %shared : !struct.type<@Holder>, !struct.type<@TPair::@Pair<[#id, #id]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: poly.template @"TPair_\1A_0" { +// CHECK-NEXT: poly.param @A : index +// CHECK-NEXT: struct.def @Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@"TPair_\1A_0"::@Pair<[@A]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@"TPair_\1A_0"::@Pair<[@A]>> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@"TPair_\1A_0"::@Pair<[@A]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@"TPair_\1A_0"::@Pair<[@A]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @"TPair_0_\1A" { +// CHECK-NEXT: poly.param @B : index +// CHECK-NEXT: struct.def @Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@"TPair_0_\1A"::@Pair<[@B]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@"TPair_0_\1A"::@Pair<[@B]>> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@"TPair_0_\1A"::@Pair<[@B]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@"TPair_0_\1A"::@Pair<[@B]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TPair_0_0_Pair { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TPair_0_0_Pair> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TPair_0_0_Pair> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TPair_0_0_Pair>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @left : !struct.type<@TPair_0_0_Pair> {column} +// CHECK-NEXT: struct.member @right : !struct.type<@TPair_0_0_Pair> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: struct.writem %[[VAL_6]][@left] = %[[VAL_7]] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: struct.writem %[[VAL_6]][@right] = %[[VAL_7]] : <@Holder>, !struct.type<@TPair_0_0_Pair> +// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_8:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_member_expandable_fail.llzk b/test/Transforms/Flattening/scalarize_member_expandable_fail.llzk new file mode 100644 index 0000000000..ff94c67942 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_expandable_fail.llzk @@ -0,0 +1,529 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + // Regression: reading an unwritten expandable index creates a result-typed pending + // materialization, which must be checked against the later split-member scalar type. + // expected-error@+1 {{cannot split heterogeneous array member because an expandable array reuses one SSA value for incompatible scalar member types}} + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{unwritten index [0 : index] is materialized for read result type '!struct.type<@TCell_1_Cell>'}} + %observed = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@wrong] = %observed : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{same index is also materialized for scalar member type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // Regression: an initialized expandable index must reject an incompatible static read + // before rewriteExpandableLocalArray starts mutating the local array. + // expected-error@+1 {{cannot split heterogeneous array member because an expandable array reuses one SSA value for incompatible scalar member types}} + %fallback = array.new %init0, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{array index [0 : index] is materialized for read result type '!struct.type<@TCell_1_Cell>'}} + %observed = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@wrong] = %observed : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{same index is also materialized for scalar member type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @wrong : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + // expected-error@+1 {{cannot split heterogeneous array member because an expandable array stores a scalar value with an incompatible read result type}} + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{array index [0 : index] stores value type '!struct.type<@TCell_0_Cell>'}} + %written0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %fallback[%c0] = %written0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + // expected-note@+1 {{same index is read as '!struct.type<@TCell_1_Cell>'}} + %observed = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@wrong] = %observed : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + // expected-error@+1 {{cannot split heterogeneous array member because expandable whole-array write uses different index set}} + struct.member @items : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %local = array.new{(%c2)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{candidate establishes 2 split member indices}} + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c3 = arith.constant 3 : index + %fallback = array.new{(%c3)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{expandable whole-array write has 3 array indices, including extra index [2 : index]}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + // expected-error@+1 {{cannot split heterogeneous array member because expandable whole-array write uses different index set}} + struct.member @items : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %local = array.new{(%c3)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %local[%c0] = %item0 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c2] = %item2 : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + // expected-note@+1 {{candidate establishes 3 split member indices}} + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c2 = arith.constant 2 : index + %fallback = array.new{(%c2)} : !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{expandable whole-array write has 2 array indices, missing index [2 : index]}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<#id x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + // expected-error@+1 {{cannot split heterogeneous array member because candidate writes require incompatible scalar member types}} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %left0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %left1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %right = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %right0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %right1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + // expected-note@+1 {{candidate writes index [0 : index] with type '!struct.type<@TCell_1_Cell>'}} + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %fallback = array.new %init0, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{conflicting candidate writes the same index with type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + // expected-error@+1 {{cannot split heterogeneous array member because an expandable array stores a scalar value with an incompatible split-member type}} + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{array index [0 : index] stores value type '!struct.type<@TCell_1_Cell>'}} + %wrong0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %ok1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %fallback[%c0] = %wrong0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %fallback[%c1] = %ok1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + // expected-note@+1 {{same index is written with scalar member type '!struct.type<@TCell_0_Cell>'}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %left0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %left1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %right = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %right0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %right1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + // expected-error@+1 {{cannot split heterogeneous array member because an expandable array reuses one SSA value for incompatible scalar member types}} + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // It's not deterministic which index is materialized first, so we have to accept either in the notes. + // expected-note-re@+1 {{unwritten index [{{[01]}} : index] is materialized for scalar member type '!struct.type<@TCell_{{[01]}}_Cell>'}} + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note-re@+1 {{same index is also materialized for scalar member type '!struct.type<@TCell_{{[01]}}_Cell>'}} + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TPair { + poly.param @A : index + poly.param @B : index + + struct.def @Pair { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TPair::@Pair<[@A, @B]>> { + %self = struct.new : !struct.type<@TPair::@Pair<[@A, @B]>> + function.return %self : !struct.type<@TPair::@Pair<[@A, @B]>> + } + + function.def @constrain(%self: !struct.type<@TPair::@Pair<[@A, @B]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @seed : !struct.type<@TPair::@Pair<[#id, #id]>> {column} + struct.member @left : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + %left0 = llzk.nondet : !struct.type<@TPair::@Pair<[#id, 0]>> + %left1 = llzk.nondet : !struct.type<@TPair::@Pair<[1, 1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[#id, 0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[1, 1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + + %right = array.new : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + %right0 = llzk.nondet : !struct.type<@TPair::@Pair<[0, #id]>> + %right1 = llzk.nondet : !struct.type<@TPair::@Pair<[1, 1]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[0, #id]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[1, 1]>> + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + } else { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + // expected-note@+1 {{array index [0 : index] stores value type '!struct.type<@TPair::@Pair<[affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>]>>'}} + %shared = struct.readm %self[@seed] : !struct.type<@Holder>, !struct.type<@TPair::@Pair<[#id, #id]>> + // expected-error@+1 {{cannot split heterogeneous array member because a scalarization candidate uses a non-nondeterministic SSA value for a refined scalar member type}} + %fallback = array.new : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + array.write %fallback[%c0] = %shared : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[#id, #id]>> + %fallback1 = llzk.nondet : !struct.type<@TPair::@Pair<[1, 1]>> + array.write %fallback[%c1] = %fallback1 : !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>>, !struct.type<@TPair::@Pair<[1, 1]>> + // expected-note@+1 {{same index is written with scalar member type '!struct.type<@"TPair_\1A_0"::@Pair<[affine_map<(d0) -> (d0)>]>>'}} + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TPair::@Pair<[#id, #id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/scalarize_member_initialized.llzk b/test/Transforms/Flattening/scalarize_member_initialized.llzk new file mode 100644 index 0000000000..73ae18b3ab --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_initialized.llzk @@ -0,0 +1,555 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %local = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %local[%c0] = %item0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c2] = %item2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c2 = arith.constant 2 : index + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %updated2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %fallback = array.new %init0, %init1, %init0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %fallback[%c2] = %updated2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_2_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column, llzk.pub} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_2_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_2_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_2_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @items_2 : !struct.type<@TCell_2_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_6:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_6]] { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TCell_2_Cell::@compute() : () -> !struct.type<@TCell_2_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_2] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @TCell_2_Cell::@compute() : () -> !struct.type<@TCell_2_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_0] = %[[VAL_13]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_1] = %[[VAL_12]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_7]][@items_2] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_7]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_14:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_15:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @observed : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %local = array.new %init0, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %read0 = array.read %local[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@observed] = %read0 : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @observed : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-DAG: %[[VAL_7:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-DAG: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-DAG: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_6]][@observed] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_6]][@items_0] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_6]][@items_1] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %left0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %left1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + %right = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %right0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %right1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %fallback = array.new %init0, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @left_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @left_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @right_2 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @right_3 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-DAG: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-DAG: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_13:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> {product_source = "init0"} + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %fallback = array.new %init0, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %left = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %left0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %left1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %left[%c0] = %left0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %left[%c1] = %left1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %left : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %right = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %right0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %right1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %right[%c0] = %right0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %right[%c1] = %right1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@right] = %right : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @left_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @left_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @right_2 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @right_3 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-DAG: %[[VAL_6:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> {product_source = "init0"} +// CHECK-DAG: %[[VAL_7:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_2] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_3] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_12:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_13:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @first : !struct.type<@TCell::@Cell<[0]>> {column} + struct.member @middle : !struct.type<@TCell::@Cell<[1]>> {column} + struct.member @last : !struct.type<@TCell::@Cell<[2]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %init0 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init2 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %updated0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %updated2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + %local = array.new %init0, %init1, %init2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %local[%c0] = %updated0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c2] = %updated2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + %first = array.read %local[%c0] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %middle = array.read %local[%c1] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %last = array.read %local[%c2] : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@first] = %first : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@middle] = %middle : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@last] = %last : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[2]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_2_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_2_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_2_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@TCell_2_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_5:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_2_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_6]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_7:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @first : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @middle : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @last : !struct.type<@TCell_2_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-NEXT: %[[VAL_13:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_14:[0-9a-zA-Z_\.]+]] = function.call @TCell_2_Cell::@compute() : () -> !struct.type<@TCell_2_Cell> +// CHECK-NEXT: struct.writem %[[VAL_8]][@first] = %[[VAL_13]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_8]][@middle] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_8]][@last] = %[[VAL_14]] : <@Holder>, !struct.type<@TCell_2_Cell> +// CHECK-NEXT: function.return %[[VAL_8]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_15:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_member_initialized_write_fail.llzk b/test/Transforms/Flattening/scalarize_member_initialized_write_fail.llzk new file mode 100644 index 0000000000..e2cd482bf9 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_initialized_write_fail.llzk @@ -0,0 +1,54 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column, llzk.pub} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %local = array.new : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %item2 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[2]>> + array.write %local[%c0] = %item0 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c2] = %item2 : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[2]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %init = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{cannot split heterogeneous array member because a scalarization candidate reuses one SSA value for incompatible scalar member types}} + %fallback = array.new %init, %init, %init : !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+2 {{value is used for scalar member type}} + // expected-note@+1 {{same value is also used for scalar member type}} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<3 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/scalarize_member_shared_external_use.llzk b/test/Transforms/Flattening/scalarize_member_shared_external_use.llzk new file mode 100644 index 0000000000..373e2970f5 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_shared_external_use.llzk @@ -0,0 +1,117 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @second : !struct.type<@TCell::@Cell<[0]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %c0 = arith.constant 0 : index + %shared = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %fallback = array.new %shared, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // The local read remains generic; the sibling candidate above supplies the concrete + // split-member requirement for this index. + %read = array.read %fallback[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@second] = %read : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK: #[[$ATTR_0:[0-9a-zA-Z_\.]+]] = affine_map<(d0) -> (d0)> +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: poly.template @TCell { +// CHECK-NEXT: poly.param @N : index +// CHECK-NEXT: struct.def @Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_GENERIC_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell::@Cell<[@N]>> +// CHECK-NEXT: function.return %[[VAL_GENERIC_0]] : !struct.type<@TCell::@Cell<[@N]>> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_GENERIC_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell::@Cell<[@N]>>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @second : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-DAG: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell::@Cell<[#[[$ATTR_0]]]>> +// CHECK-DAG: %[[VAL_10:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> +// CHECK-DAG: %[[VAL_11:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@second] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_0] = %[[VAL_10]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@items_1] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_10:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_11:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_member_shared_external_use_fail.llzk b/test/Transforms/Flattening/scalarize_member_shared_external_use_fail.llzk new file mode 100644 index 0000000000..9bbb79407b --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_shared_external_use_fail.llzk @@ -0,0 +1,211 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column, llzk.pub} // expected-note {{split member defined here}} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } + + function.def @bad(%holder: !struct.type<@Holder>) -> !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> { + // expected-error@+1 {{cannot split heterogeneous array member because a whole-array read has an unsupported use}} + %items = struct.readm %holder[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> // expected-note {{unsupported use is here}} + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column, llzk.pub} // expected-note {{split member defined here}} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } + + struct.def @Consumer { + function.def @compute(%holder: !struct.type<@Holder>, %idx: index) -> !struct.type<@Consumer> { + %self = struct.new : !struct.type<@Consumer> + %items = struct.readm %holder[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-error@+1 {{cannot split heterogeneous array member because a read of it uses a dynamic array index}} + %item = array.read %items[%idx] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + function.return %self : !struct.type<@Consumer> + } + + function.def @constrain(%self: !struct.type<@Consumer>, %holder: !struct.type<@Holder>, %idx: index) { + function.return + } + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column, llzk.pub} // expected-note {{split member defined here}} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } + + function.def @bad(%holder: !struct.type<@Holder>) { + %c0 = arith.constant 0 : index + %items = struct.readm %holder[@items] : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-error@+1 {{cannot split heterogeneous array member because a read result type is incompatible with the split scalar member type}} + %item = array.read %items[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + function.return + } +} + +// ----- + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @second : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + scf.if %cond { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %local = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %local[%c0] = %item0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %local[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %local : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %shared = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %init1 = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{cannot split heterogeneous array member because a scalarization candidate reuses one SSA value for incompatible scalar member types}} + %fallback = array.new %shared, %init1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+2 {{value is used for member type}} + // expected-note@+2 {{same value is also used for array index}} + struct.writem %self[@second] = %shared : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} diff --git a/test/Transforms/Flattening/scalarize_member_shared_nondet_no_read_external_constraint.llzk b/test/Transforms/Flattening/scalarize_member_shared_nondet_no_read_external_constraint.llzk new file mode 100644 index 0000000000..ac7363e9ea --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_shared_nondet_no_read_external_constraint.llzk @@ -0,0 +1,88 @@ +// RUN: llzk-opt -llzk-flatten %s | FileCheck %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @items : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @product() -> !struct.type<@Holder> + attributes {function.allow_witness, function.allow_constraint} { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + // Establish the concrete split-member types. + %first0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %first1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %first = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %first[%c0] = %first0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %first[%c1] = %first1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@items] = %first : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + + // No local read refines %shared. Its generic constraint and scalar member write must still + // use the same concrete nondeterministic witness after scalarization. + %shared = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> {product_source = "shared"} + %other = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %fallback = array.new %shared, %other : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.call @TCell::@Cell::@constrain(%shared) : (!struct.type<@TCell::@Cell<[#id]>>) -> () {product_source = "shared"} + struct.writem %self[@items] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @items_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @product() -> !struct.type<@Holder> attributes {function.allow_constraint, function.allow_witness} { +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_5]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_6]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-DAG: %[[VAL_7:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_1_Cell> +// CHECK-DAG: %[[VAL_8:[0-9a-zA-Z_\.]+]] = llzk.nondet : !struct.type<@TCell_0_Cell> {product_source = "shared"} +// CHECK-DAG: function.call @TCell_0_Cell::@constrain(%[[VAL_8]]) : (!struct.type<@TCell_0_Cell>) -> () {product_source = "shared"} +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_4]][@items_1] = %[[VAL_7]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_4]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_member_shared_unifiable_cast.llzk b/test/Transforms/Flattening/scalarize_member_shared_unifiable_cast.llzk new file mode 100644 index 0000000000..d9efbd9d1b --- /dev/null +++ b/test/Transforms/Flattening/scalarize_member_shared_unifiable_cast.llzk @@ -0,0 +1,111 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s | FileCheck --enable-var-scope %s + +// Regression test for an expandable array whose generic unifiable-cast value is reused by +// multiple whole-array writes. Scalarization must give both generated scalar writes the shared +// concrete cast value rather than leaving both writes attached to the generic cast result. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @left : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + struct.member @right : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute(%cond: i1) -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.if %cond { + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %genericItem0 = poly.unifiable_cast %item0 : (!struct.type<@TCell::@Cell<[0]>>) -> !struct.type<@TCell::@Cell<[#id]>> + %genericItem1 = poly.unifiable_cast %item1 : (!struct.type<@TCell::@Cell<[1]>>) -> !struct.type<@TCell::@Cell<[#id]>> + %candidate = array.new %genericItem0, %genericItem1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@left] = %candidate : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@right] = %candidate : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } else { + %fallback = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + %item0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %shared = poly.unifiable_cast %item0 : (!struct.type<@TCell::@Cell<[0]>>) -> !struct.type<@TCell::@Cell<[#id]>> + %item1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + array.write %fallback[%c0] = %shared : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[#id]>> + array.write %fallback[%c1] = %item1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + struct.writem %self[@left] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@right] = %fallback : !struct.type<@Holder>, !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + } + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>, %cond: i1) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @TCell_1_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_1_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_1_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_3:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_1_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @left_2 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @left_3 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: struct.member @right_0 : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @right_1 : !struct.type<@TCell_1_Cell> {column} +// CHECK-NEXT: function.def @compute(%[[VAL_4:[0-9a-zA-Z_\.]+]]: i1) -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: scf.if %[[VAL_4]] { +// CHECK-NEXT: %[[VAL_6:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_7:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: %[[VAL_8:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_6]] : (!struct.type<@TCell_0_Cell>) -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_9:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_7]] : (!struct.type<@TCell_1_Cell>) -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_2] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_3] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_0] = %[[VAL_8]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_1] = %[[VAL_9]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } else { +// CHECK-NEXT: %[[VAL_10:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_11:[0-9a-zA-Z_\.]+]] = poly.unifiable_cast %[[VAL_10]] : (!struct.type<@TCell_0_Cell>) -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_12:[0-9a-zA-Z_\.]+]] = function.call @TCell_1_Cell::@compute() : () -> !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_2] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@left_3] = %[[VAL_12]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_0] = %[[VAL_11]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_5]][@right_1] = %[[VAL_12]] : <@Holder>, !struct.type<@TCell_1_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.return %[[VAL_5]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_13:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>, %[[VAL_14:[0-9a-zA-Z_\.]+]]: i1) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/scalarize_shared_nondet_external_use_preflight_fail.llzk b/test/Transforms/Flattening/scalarize_shared_nondet_external_use_preflight_fail.llzk new file mode 100644 index 0000000000..c1b78a0351 --- /dev/null +++ b/test/Transforms/Flattening/scalarize_shared_nondet_external_use_preflight_fail.llzk @@ -0,0 +1,59 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + function.def @consume0(%value: !struct.type<@TCell::@Cell<[0]>>) { + function.return + } + + function.def @consume1(%value: !struct.type<@TCell::@Cell<[1]>>) { + function.return + } + + function.def @consumeGeneric(%value: !struct.type<@TCell::@Cell<[#id]>>) { + function.return + } + + function.def @preflight_external_use() attributes {function.allow_witness} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + + // This independent candidate is ordered before the failing candidate. + %first0 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + %first1 = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[1]>> + %first = array.new : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + array.write %first[%c0] = %first0 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + array.write %first[%c1] = %first1 : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + %firstRead0 = array.read %first[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + %firstRead1 = array.read %first[%c1] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[1]>> + function.call @consume0(%firstRead0) : (!struct.type<@TCell::@Cell<[0]>>) -> () + function.call @consume1(%firstRead1) : (!struct.type<@TCell::@Cell<[1]>>) -> () + + %generic = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + %other = llzk.nondet : !struct.type<@TCell::@Cell<[#id]>> + // expected-error@+1 {{cannot scalarize array because a generic nondeterministic initializer has an unsupported external use}} + %later = array.new %generic, %other : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>> + // expected-note@+1 {{source witness is also used here}} + function.call @consumeGeneric(%generic) : (!struct.type<@TCell::@Cell<[#id]>>) -> () + %laterRead = array.read %later[%c0] : !array.type<2 x !struct.type<@TCell::@Cell<[#id]>>>, !struct.type<@TCell::@Cell<[0]>> + function.call @consume0(%laterRead) : (!struct.type<@TCell::@Cell<[0]>>) -> () + function.return + } +} diff --git a/test/Transforms/Flattening/test/Transforms/Flattening/member_read_array_initializer_reproducer.llzk b/test/Transforms/Flattening/test/Transforms/Flattening/member_read_array_initializer_reproducer.llzk new file mode 100644 index 0000000000..10401e94f9 --- /dev/null +++ b/test/Transforms/Flattening/test/Transforms/Flattening/member_read_array_initializer_reproducer.llzk @@ -0,0 +1,136 @@ +// RUN: llzk-opt -split-input-file -llzk-flatten -verify-diagnostics %s | FileCheck %s + +// A concrete member write refines the generic member read to Cell<0>. array.new must follow that +// refinement so its initializer and declared element type remain identical. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @item : !struct.type<@TCell::@Cell<[#id]>> {column} + struct.member @items : !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@item] = %item : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + %read = struct.readm %self[@item] : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + %items = array.new %read : !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> + struct.writem %self[@items] = %items : !struct.type<@Holder>, !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @item : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: struct.member @items : !array.type<1 x !struct.type<@TCell_0_Cell>> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_2]][@item] = %[[VAL_3]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_4:[0-9a-zA-Z_\.]+]] = struct.readm %[[VAL_2]][@item] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.new %[[VAL_4]] : <1 x !struct.type<@TCell_0_Cell>> +// CHECK-NEXT: struct.writem %[[VAL_2]][@items] = %[[VAL_5]] : <@Holder>, !array.type<1 x !struct.type<@TCell_0_Cell>> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_6:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// This input verifies before flattening: the generic member read exactly matches the generic +// element type required by array.new. A concrete member write then refines the member and its read +// to Cell<0>, but array.new retains Cell<#id> as its element type and fails final verification. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @item : !struct.type<@TCell::@Cell<[#id]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + struct.writem %self[@item] = %item : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[0]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + %item = struct.readm %self[@item] : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + %items = array.new %item : !array.type<1 x !struct.type<@TCell::@Cell<[#id]>>> + function.return + } + } +} +// CHECK-LABEL: module attributes {llzk.lang} { +// CHECK-NEXT: struct.def @TCell_0_Cell { +// CHECK-NEXT: struct.member @value : !felt.type {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@TCell_0_Cell> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_0:[0-9a-zA-Z_\.]+]] = struct.new : <@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_0]] : !struct.type<@TCell_0_Cell> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_1:[0-9a-zA-Z_\.]+]]: !struct.type<@TCell_0_Cell>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: struct.def @Holder { +// CHECK-NEXT: struct.member @item : !struct.type<@TCell_0_Cell> {column} +// CHECK-NEXT: function.def @compute() -> !struct.type<@Holder> attributes {function.allow_witness} { +// CHECK-NEXT: %[[VAL_2:[0-9a-zA-Z_\.]+]] = struct.new : <@Holder> +// CHECK-NEXT: %[[VAL_3:[0-9a-zA-Z_\.]+]] = function.call @TCell_0_Cell::@compute() : () -> !struct.type<@TCell_0_Cell> +// CHECK-NEXT: struct.writem %[[VAL_2]][@item] = %[[VAL_3]] : <@Holder>, !struct.type<@TCell_0_Cell> +// CHECK-NEXT: function.return %[[VAL_2]] : !struct.type<@Holder> +// CHECK-NEXT: } +// CHECK-NEXT: function.def @constrain(%[[VAL_4:[0-9a-zA-Z_\.]+]]: !struct.type<@Holder>) attributes {function.allow_constraint} { +// CHECK-NEXT: function.return +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } diff --git a/test/Transforms/Flattening/unifiable_cast_member_write_type_guard.llzk b/test/Transforms/Flattening/unifiable_cast_member_write_type_guard.llzk new file mode 100644 index 0000000000..5594a3f329 --- /dev/null +++ b/test/Transforms/Flattening/unifiable_cast_member_write_type_guard.llzk @@ -0,0 +1,46 @@ +// RUN: llzk-opt -llzk-flatten -verify-diagnostics %s + +// A generic unifiable cast may feed member writes that are independently instantiated. Its input +// is Cell<0>, which satisfies @zero but not @one, so flattening must not retag the shared cast +// result to Cell<0>. +#id = affine_map<(i) -> (i)> +module attributes {llzk.lang} { + poly.template @TCell { + poly.param @N : index + + // expected-warning@+1 {{Parameterized definition still has uses!}} + struct.def @Cell { + struct.member @value : !felt.type {column} + + function.def @compute() -> !struct.type<@TCell::@Cell<[@N]>> { + %self = struct.new : !struct.type<@TCell::@Cell<[@N]>> + function.return %self : !struct.type<@TCell::@Cell<[@N]>> + } + + function.def @constrain(%self: !struct.type<@TCell::@Cell<[@N]>>) { + function.return + } + } + } + + struct.def @Holder { + struct.member @zero : !struct.type<@TCell::@Cell<[0]>> {column} + struct.member @one : !struct.type<@TCell::@Cell<[1]>> {column} + + function.def @compute() -> !struct.type<@Holder> { + %self = struct.new : !struct.type<@Holder> + %item = function.call @TCell::@Cell::@compute() : () -> !struct.type<@TCell::@Cell<[0]>> + // The cast cannot be finalized because its shared result has incompatible member targets. + // It must not instead be retagged to Cell<0>, which would make the @one write invalid. + // expected-error@+1 {{'poly.unifiable_cast' op input type '!struct.type<@TCell_0_Cell>' and output type '!struct.type<@TCell::@Cell<[affine_map<(d0) -> (d0)>]>>' are not unifiable}} + %generic = poly.unifiable_cast %item : (!struct.type<@TCell::@Cell<[0]>>) -> !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@zero] = %generic : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + struct.writem %self[@one] = %generic : !struct.type<@Holder>, !struct.type<@TCell::@Cell<[#id]>> + function.return %self : !struct.type<@Holder> + } + + function.def @constrain(%self: !struct.type<@Holder>) { + function.return + } + } +}