Ops.td

Go to the documentation of this file.

//===-- Ops.td ---------------------------------------------*- tablegen -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
// Adapted from mlir/include/mlir/Dialect/Func/IR/FuncOps.td
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#ifndef LLZK_STRUCT_OPS
#define LLZK_STRUCT_OPS
 
include "llzk/Dialect/Function/IR/OpTraits.td"
include "llzk/Dialect/Struct/IR/Dialect.td"
include "llzk/Dialect/Struct/IR/OpInterfaces.td"
include "llzk/Dialect/Struct/IR/Types.td"
include "llzk/Dialect/Shared/OpTraits.td"
include "llzk/Dialect/Verif/IR/OpInterfaces.td"
 
include "mlir/IR/OpAsmInterface.td"
include "mlir/IR/RegionKindInterface.td"
include "mlir/IR/SymbolInterfaces.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
 
class StructDialectOp<string mnemonic, list<Trait> traits = []>
    : Op<StructDialect, mnemonic, traits>;
 
/// Only valid/implemented for StructDefOp. Sets the proper `AllowWitnessAttr`
/// and `AllowConstraintAttr` on the functions defined within the StructDefOp.
def SetFuncAllowAttrs : NativeOpTrait<"SetFuncAllowAttrs">, StructuralOpTrait {
  string cppNamespace = "::llzk::component";
}
 
//===------------------------------------------------------------------===//
// Struct Operations
//===------------------------------------------------------------------===//
 
def LLZK_StructDefOp
    : StructDialectOp<"def", [ParentOneOf<["::mlir::ModuleOp",
                                           "::llzk::polymorphic::TemplateOp"]>,
                              Symbol, LLZKSymbolTable, IsolatedFromAbove,
                              SetFuncAllowAttrs, NoRegionArguments,
                              ContractTarget]#GraphRegionNoTerminator.traits> {
  let summary = "circuit component definition";
  let description = [{
      This operation describes a component in a circuit. It can contain any number
      of members that hold inputs, outputs, intermediate values, and subcomponents
      of the defined component. It also contains a `compute()` function that holds
      the witness generation code for the component and a `constrain()` function
      that holds that constraint generation code for the component.
 
      Example:
 
      ```llzk
      struct.def @ComponentA {
        member @f1 : !array.type<5 x index>
        member @f2 : !felt.type {llzk.pub}
 
        function.def @compute(%p: !felt.type) -> !struct.type<@ComponentA> {
          %self = struct.new : !struct.type<@ComponentA>
          // initialize all members of `%self` here
          return %self : !struct.type<@ComponentA>
        }
 
        function.def @constrain(%self: !struct.type<@ComponentA>, %p: !felt.type) {
          // emit constraints here
          return
        }
      }
      ```
 
      The optional `llzk.main = !struct.type<...>` attribute on the top-level module in an
      LLZK IR program expresses the main entry point of the circuit as a concrete instantiation
      of a specific struct. For example, `llzk.main = !struct.type<@Top<[52,12]>>` or
      `llzk.main = !struct.type<@Main<[i1, !felt.type, 256]>>`.
      This struct has additional restrictions:
      1. The parameter types of its functions (besides the required "self" parameter) can
         only be `!felt.type` or `!array.type<.. x !felt.type>`.
      2. All inputs to the main struct are signals, even though they are not marked with
         the `{signal}` attribute (as the `{signal}` attribute is not used on function arguments).
         Input signals may be public (given the `{llzk.pub}` attribute) or private (by default).
      3. All public members of the main struct are also signals (representing public output signals),
         and can also only be `!felt.type` or `!array.type<.. x !felt.type>`.
 
      Example of a `Main` component:
 
      ```llzk
      module attributes {llzk.main = !struct.type<@Main>, llzk.lang} {
        struct.def @Main {
          member @out : !felt.type {llzk.pub, signal} // public output signal
          member @out2 : !felt.type {llzk.pub} // also a public output signal, implicitly
          member @intermediate : !struct.type<@OtherStruct> // an intermediate value, not a signal, can be any LLZK type
          // %p is a private input signal
          function.def @compute(%p: !felt.type) -> !struct.type<@Main> {
            // ...
          }
          // %self is not a signal, but again, %p is a private input signal
          function.def @constrain(%self: !struct.type<@Main>, %p: !felt.type) {
            // ...
          }
        }
      }
      ```
  }];
 
  let arguments = (ins SymbolNameAttr:$sym_name);
 
  let regions = (region SizedRegion<1>:$bodyRegion);
 
  let assemblyFormat = [{ $sym_name $bodyRegion attr-dict }];
 
  let useCustomPropertiesEncoding = 1;
 
  let extraClassDeclaration = [{
    /// Gets the StructType representing this struct. If the `constParams` to use in
    /// the type are not given, the StructType will use `this->getParams()`.
    StructType getType(::std::optional<::mlir::ArrayAttr> constParams = {});
    inline StructType getType(::std::optional<::mlir::ArrayAttr> constParams = {}) const {
      return const_cast<StructDefOp*>(this)->getType(constParams);
    }
 
    /// Gets the MemberDefOp that defines the member in this
    /// structure with the given name, if present.
    MemberDefOp getMemberDef(::mlir::StringAttr memberName);
 
    /// Get all MemberDefOp in this structure.
    ::std::vector<MemberDefOp> getMemberDefs();
 
    /// Returns whether the struct defines members marked as columns.
    ::mlir::LogicalResult hasColumns() {
      return ::mlir::success(::llvm::any_of(getMemberDefs(), [](MemberDefOp memberOp) {
        return memberOp.getColumn();
      }));
    }
 
    /// Returns whether the struct defines members marked as signals.
    ::mlir::LogicalResult hasSignals() {
      return ::mlir::success(::llvm::any_of(getMemberDefs(), [](MemberDefOp memberOp) {
        return memberOp.getSignal();
      }));
    }
 
    /// Gets the FuncDefOp that defines the compute function in this structure, if present, or `nullptr` otherwise.
    ::llzk::function::FuncDefOp getComputeFuncOp();
 
    /// Gets the FuncDefOp that defines the constrain function in this structure, if present, or `nullptr` otherwise.
    ::llzk::function::FuncDefOp getConstrainFuncOp();
 
    /// Gets the FuncDefOp that defines the product function in this structure, if present, or `nullptr` otherwise
    ::llzk::function::FuncDefOp getProductFuncOp();
 
    /// Returns `true` iff this structure defines compute and constrain functions.
    bool hasComputeConstrain() { return lookupSymbol(FUNC_NAME_COMPUTE) != nullptr && lookupSymbol(FUNC_NAME_CONSTRAIN) != nullptr; }
 
    /// Generate header string, in the same format as the assemblyFormat.
    ::std::string getHeaderString();
 
    /// Return `true` iff the `struct.def` appears within a `poly.template` that defines
    /// constant parameters and/or expressions.
    bool hasTemplateSymbolBindings();
 
    /// If this `struct.def` is within a `poly.template`, return names of all `poly.param`
    /// within the `poly.template` in the order they are defined. Otherwise, return empty.
    /// The names are returned as `FlatSymbolRefAttr` but the more general `Attribute` type
    /// is used in the return type since that's usually what's needed.
    ::llvm::SmallVector<::mlir::Attribute> getTemplateParamOpNames();
 
    /// If this `struct.def` is within a `poly.template`, return names of all `poly.expr`
    /// within the `poly.template` in the order they are defined. Otherwise, return empty.
    /// The names are returned as `FlatSymbolRefAttr` but the more general `Attribute` type
    /// is used in the return type since that's usually what's needed.
    ::llvm::SmallVector<::mlir::Attribute> getTemplateExprOpNames();
 
    /// Return the full name for this struct from the root module, including
    /// any surrounding module scopes.
    ::mlir::SymbolRefAttr getFullyQualifiedName();
 
    /// Return `true` iff this `struct.def` is the main struct. See `llzk::MAIN_ATTR_NAME`.
    bool isMainComponent();
  }];
 
  let hasRegionVerifier = 1;
}
 
def LLZK_MemberDefOp
    : StructDialectOp<
          "member", [HasParent<"::llzk::component::StructDefOp">,
                     DeclareOpInterfaceMethods<SymbolUserOpInterface>,
                     Symbol]> {
  let summary = "struct member definition";
  let description = [{
    This operation describes a member in a struct/component.
 
    Example:
 
    ```llzk
    struct.member @f1 : !felt.type
    struct.member @f2 : !felt.type {llzk.pub}
    struct.member @col1 : !felt.type {column}
    struct.member @sig1 : !felt.type {signal}
    struct.member @colsig1 : !felt.type {column, signal}
    struct.member @pubcol : !felt.type {llzk.pub, column}
    ```
 
    - Members marked with the `{llzk.pub}` attribute are considered public
      and represent outputs of their defining struct/component.
    - Members marked with the `{column}` attribute can be read/written with
      table offsets using the `readm` and `writem` operations.
    - Members marked with the `{signal}` attribute are constraint variables that
      are stored in the witness; non-signal values are intermediate expressions.
 
    Further restrictions on the `signal` attribute:
    - It is only used on struct members, since they represent storage locations.
      It is not used to mark struct inputs, as struct may be passed signal or non-signal
      (i.e., intermediate expression) values.
    - Only !felt.type or simple aggregates of !felt.type (i.e., !array.type and !pod.type)
      may be marked as signals.
  }];
 
  let arguments = (ins SymbolNameAttr:$sym_name, TypeAttrOf<AnyLLZKType>:$type,
      UnitAttr:$column, UnitAttr:$signal);
 
  // Define builders manually to avoid the default ones that have extra
  // TypeRange parameters that must always be empty.
  let skipDefaultBuilders = 1;
  let builders =
      [OpBuilder<(ins "::mlir::StringAttr":$sym_name, "::mlir::TypeAttr":$type,
           CArg<"bool", "false">:$isSignal, CArg<"bool", "false">:$isColumn)>,
       OpBuilder<(ins "::llvm::StringRef":$sym_name, "::mlir::Type":$type,
           CArg<"bool", "false">:$isSignal, CArg<"bool", "false">:$isColumn)>,
       OpBuilder<(ins "::mlir::TypeRange":$resultTypes,
           "::mlir::ValueRange":$operands,
           "::llvm::ArrayRef<::mlir::NamedAttribute>":$attributes,
           CArg<"bool", "false">:$isSignal, CArg<"bool", "false">:$isColumn)>,
       // Simpler version since 'resultTypes' and 'operands' must be empty
       OpBuilder<
           (ins "::llvm::ArrayRef<::mlir::NamedAttribute>":$attributes,
               CArg<"bool", "false">:$isSignal,
               CArg<"bool", "false">:$isColumn),
           [{ build($_builder, $_state, {}, {}, attributes, isSignal, isColumn); }]>];
 
  let assemblyFormat = [{ $sym_name `:` $type attr-dict }];
 
  let extraClassDeclaration = [{
    /// Returns whether this member is a public output.
    inline bool hasPublicAttr() { return getOperation()->hasAttr(llzk::PublicAttr::name); }
 
    /// Adds or removes the unit `llzk.pub` attribute according to `newValue`.
    void setPublicAttr(bool newValue = true);
  }];
 
  let hasVerifier = 1;
}
 
class MemberRefOpBase<string mnemonic, list<Trait> traits = []>
    : StructDialectOp<
          mnemonic, traits#[DeclareOpInterfaceMethods<MemberRefOpInterface>,
                            DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
  bit isRead = ?; // read(1) vs write(0) ops
  let extraClassDeclaration = [{
    /// Gets the definition for the `member` referenced in this op.
    inline ::mlir::FailureOr<SymbolLookupResult<MemberDefOp>> getMemberDefOp(::mlir::SymbolTableCollection &tables) {
      return ::llvm::cast<MemberRefOpInterface>(getOperation()).getMemberDefOp(tables);
    }
  }];
  let extraClassDefinition = [{
    /// Return `true` if the op is a read, `false` if it's a write.
    bool $cppClass::isRead() {
      return }]#!if(isRead, "true", "false")#[{;
    }
  }];
}
 
def LLZK_MemberReadOp
    : MemberRefOpBase<"readm", [VerifySizesForMultiAffineOps<1>,
                                MemoryEffects<[MemRead]>]> {
  let summary = "read value of a struct member";
  let description = [{
    This operation copies the value of a named member in a struct/component.
    The result is a value copy, not a pointer or alias to the member's storage.
    When the member itself has `struct.type`, the copied value is an
    identity-bearing handle to the same component instance.
 
    A struct can read its own members regardless of whether they are marked as
    public (i.e., with the `llzk.pub` attribute) or private (members without the
    `llzk.pub` attribute). However, when reading members of other components, only
    public members can be accessed. Free functions may also only read public members.
 
    The value can be read from the signals table, in which case it can be
    offset by a constant value. A negative value represents reading a value
    backwards and a positive value represents reading a value forward.
    Only members marked as columns can be read in this manner.
  }];
  let isRead = 1;
 
  // See `VerifySizesForMultiAffineOps` for more explanation of these arguments.
  let arguments = (ins LLZK_StructType:$component,
      FlatSymbolRefAttr:$member_name,
      OptionalAttr<AnyAttrOf<[SymbolRefAttr, IndexAttr,
                              AffineMapAttr]>>:$tableOffset,
      // List of AffineMap operand groups where each group provides the
      // arguments to instantiate the next (left-to-right) AffineMap used in
      // `tableOffset`.
      VariadicOfVariadic<Index, "mapOpGroupSizes">:$mapOperands,
      // Within each group in '$mapOperands', denotes the number of values that
      // are AffineMap "dimensional" arguments with the remaining values being
      // AffineMap "symbolic" arguments.
      DefaultValuedAttr<DenseI32ArrayAttr, "{}">:$numDimsPerMap,
      // Denotes the size of each variadic group in '$mapOperands'.
      DenseI32ArrayAttr:$mapOpGroupSizes);
  let results = (outs AnyLLZKType:$val);
 
  // Define builders manually so inference of operand layout attributes is not
  // circumvented.
  let skipDefaultBuilders = 1;
  let builders =
      [OpBuilder<(ins "::mlir::Type":$resultType, "::mlir::Value":$component,
           "::mlir::StringAttr":$member)>,
       OpBuilder<(ins "::mlir::Type":$resultType, "::mlir::Value":$component,
           "::mlir::StringAttr":$member, "::mlir::Attribute":$dist,
           "::mlir::ValueRange":$mapOperands,
           "std::optional<int32_t>":$numDims)>,
       OpBuilder<(ins "::mlir::Type":$resultType, "::mlir::Value":$component,
                     "::mlir::StringAttr":$member,
                     "::mlir::SymbolRefAttr":$dist),
                 [{
              build($_builder, $_state, resultType, component, member, dist, ::mlir::ValueRange(), std::nullopt);
           }]>,
       OpBuilder<(ins "::mlir::Type":$resultType, "::mlir::Value":$component,
                     "::mlir::StringAttr":$member, "::mlir::IntegerAttr":$dist),
                 [{
              build($_builder, $_state, resultType, component, member, dist, ::mlir::ValueRange(), std::nullopt);
           }]>,
       OpBuilder<(ins "::mlir::TypeRange":$resultTypes,
           "::mlir::ValueRange":$operands,
           "::mlir::ArrayRef<::mlir::NamedAttribute>":$attrs)>];
 
  let assemblyFormat = [{
    $component `[` $member_name `]`
    ( `{` custom<MultiDimAndSymbolList>($mapOperands, $numDimsPerMap)^ `}` )?
    `:` type($component) `,` type($val)
    attr-dict
  }];
 
  let hasVerifier = 1;
}
 
def LLZK_MemberWriteOp
    : MemberRefOpBase<"writem", [WitnessGen, MemoryEffects<[MemWrite]>]> {
  let summary = "write value to a struct member";
  let description = [{
    This operation copies a value to a named member in a struct/component.
    It does not create a pointer or alias between the source value and the member's
    storage. When the value has `struct.type`, the copied value is an
    identity-bearing handle to the same component instance.
 
    A struct can write its own members. However, writing members of other components
    is not allowed.
  }];
  let isRead = 0;
 
  let arguments = (ins LLZK_StructType:$component,
      FlatSymbolRefAttr:$member_name, AnyLLZKType:$val);
 
  let assemblyFormat = [{
    $component `[` $member_name `]` `=` $val `:` type($component) `,` type($val) attr-dict
  }];
}
 
def LLZK_CreateStructOp
    : StructDialectOp<"new", [DeclareOpInterfaceMethods<
                                  OpAsmOpInterface, ["getAsmResultNames"]>,
                              DeclareOpInterfaceMethods<SymbolUserOpInterface>,
                              WitnessGen,
]> {
  let summary = "create a new struct";
  let description = [{
    This operation creates a new, uninitialized instance of a struct.
 
    Example:
 
    ```llzk
    %self = struct.new : !struct.type<@Reg>
    ```
  }];
 
  let results = (outs LLZK_StructType:$result);
 
  let assemblyFormat = [{ `:` type($result) attr-dict }];
}
 
#endif // LLZK_STRUCT_OPS