LLZK 0.1.0
Veridise's ZK Language IR
Loading...
Searching...
No Matches
OpTraits.td
Go to the documentation of this file.
1//===-- OpTraits.td - Custom Trait classes for ops ---------*- tablegen -*-===//
2//
3// Part of the LLZK Project, under the Apache License v2.0.
4// See LICENSE.txt for license information.
5// Copyright 2025 Veridise Inc.
6// SPDX-License-Identifier: Apache-2.0
7//
8//===----------------------------------------------------------------------===//
9
10#ifndef LLZK_SHARED_OP_HELPER
11#define LLZK_SHARED_OP_HELPER
12
13include "mlir/Interfaces/InferTypeOpInterface.td"
14include "mlir/Interfaces/SideEffectInterfaces.td"
15
16// Implements verification for ops with an affine_map instantiation list. These
17// ops are expected to contain the following in their `arguments` list:
18// - VariadicOfVariadic<Index, "mapOpGroupSizes">:$mapOperands
19// - DefaultValuedAttr<DenseI32ArrayAttr, "{}">:$numDimsPerMap
20// - DenseI32ArrayAttr:$mapOpGroupSizes
21// Additionally, if the op also has the `AttrSizedOperandSegments` trait, the
22// parameter of this trait specifies the index within the `operandSegmentSizes`
23// attribute associated with the `$mapOperands` argument, otherwise the
24// parameter is ignored. All of these attributes are necessary because MLIR
25// stores all operands for an Op in a single list. These attributes specify how
26// the list of operands is split into logical pieces for the operand components.
27//
28// For example, suppose the `CreateArrayOp` is used to create an array with type
29// `!array.type<affine_map<(d0)->(d0)>,affine_map<(d0)[s0]->(d0+s0)> x i1>`
30//
31// 1) `CreateArrayOp` requires the `AttrSizedOperandSegments` trait because it
32// defines two variadic arguments: `$elements` and `$mapOperands` (in that
33// order). Thus, the `operandSegmentSizes` attribute is automatically defined
34// to specify the number of operands that belong to each variadic argument:
35// `operandSegmentSizes = array<i32: COUNT($elements), COUNT($mapOperands)>`
36// In the case of `CreateArrayOp`, one of those sizes will always be 0 because
37// its assembly format has `$elements` and `$mapOperands` as alternatives. In
38// this example, `COUNT($elements) = 0` and `COUNT($mapOperands) = 3` (this is
39// the sum of operand count for all affine_map that are used as array dimensions
40// in the result array type).
41//
42// 2) The `$mapOpGroupSizes` attribute groups the `$mapOperands` per affine_map.
43// This implies that their sum equals `COUNT($mapOperands)`. In the example, the
44// first affine_map has 1 parameter and the second has 2 so:
45// `mapOpGroupSizes = array<i32: 1, 2>`
46//
47// 3) Finally, the `$numDimsPerMap` attribute splits the `$mapOperands` in each
48// group into the dimensional and symbolic inputs for each affine_map.
49// Dimensional inputs appear between the () and symbolic inputs appear between
50// the []. LLZK mainly uses dimensional inputs and not symbolic inputs but both
51// are fully supported. The length of `$numDimsPerMap` must equal the length of
52// `$mapOpGroupSizes` and each element of `$numDimsPerMap` must be less than the
53// corresponding element of `$mapOpGroupSizes`. In the example, the both
54// affine_map instantiations in the array type have 1 dimensional input so:
55// `numDimsPerMap = array<i32: 1, 1>`
56//
57// It is also recomended to use `custom<AttrDictWithWarnings>(attr-dict)` in the
58// assembly format (or the associated parse/print functions directly) to parse
59// the attribute dictionary in these ops and present warnings if the
60// aforementioned attributes are manually specified.
61class VerifySizesForMultiAffineOps<int operandSegmentIndex>
62 : ParamNativeOpTrait<"VerifySizesForMultiAffineOps",
63 ""#operandSegmentIndex>,
64 StructuralOpTrait {
65 string cppNamespace = "::llzk";
66}
67
68// Identical to `TypesMatchWith` with `rhsArg = result`. This should be used
69// instead of `TypesMatchWith` when custom return type inference is necessary
70// (via `InferTypeOpAdaptor*`) because MLIR has special handing for
71// `TypesMatchWith` that results in "error: redefinition of 'inferReturnTypes'".
72class TypeMatchResultWith<string lhsArg, string lhsSummary = lhsArg,
73 string transform,
74 string comparator = "std::equal_to<>()">
75 : PredOpTrait<
76 "result type matches with "#lhsSummary#" type",
77 CPred<comparator#"("#!subst("$_self", "$"#lhsArg#".getType()",
78 transform)#", $result.getType())">> {
79 string lhs = lhsArg;
80 string rhs = "result";
81 string transformer = transform;
82}
83
84// Like TypesUnify with `rhsArg = "result"`
85class TypeUnifyWithResult<string lhsArg, string lhsSummary = lhsArg,
86 string transform = "$_self">
87 : TypeMatchResultWith<lhsArg, lhsSummary, transform, "::llzk::typesUnify">;
88
89// Implementation of TypesMatchWith for Variadic `rhsArg` that returns success
90// if `rhsArg` is empty.
91class VariadicTypesMatchWith<string summary, string lhsArg, string rhsArg,
92 string transform,
93 string comparator = "std::equal_to<>()">
94 : TypesMatchWith<
95 summary, lhsArg, rhsArg, transform,
96 "get"#snakeCaseToCamelCase<rhsArg>.ret#"().empty() || "#comparator>;
97
98// Type constraint `llzk::typesUnify(transform(lhs.getType()), rhs.getType())`.
99// If either parameter is `$result` it is recommended to use TypeUnifyWithResult
100// instead as this is likely too restrictive when type variables are involved.
101class TypesUnify<string lhsArg, string rhsArg, string lhsSummary = lhsArg,
102 string rhsSummary = rhsArg, string transform = "$_self">
103 : TypesMatchWith<rhsSummary#" type matches with "#lhsSummary#" type",
104 lhsArg, rhsArg, transform, "::llzk::typesUnify">;
105
106// Returns success if `elementArg` unifies with the `arrayArg` element type.
107class ArrayElemTypeUnifyWith<string arrayArg, string elementArg>
108 : TypesUnify<
109 arrayArg, elementArg, arrayArg#" element", elementArg,
110 "::llvm::cast<::llzk::array::ArrayType>($_self).getElementType()">;
111
112// Returns success if `$result` unifies with the `arrayArg` element type.
113class ArrayElemTypeUnifyWithResult<string arrayArg>
114 : TypeMatchResultWith<
115 arrayArg, arrayArg#" element",
116 "::llvm::cast<::llzk::array::ArrayType>($_self).getElementType()",
117 "::llzk::typesUnify">;
118
119// ArrayElemTypeUnifyWithResult + InferTypeOpAdaptorWithIsCompatible (i.e.
120// generate inferReturnTypes() and isCompatibleReturnTypes() functions)
121class ArrayTypeElemsUnifyWithResultCustomInfer<string arrayArg>
122 : TraitList<[ArrayElemTypeUnifyWithResult<arrayArg>,
123 InferTypeOpAdaptorWithIsCompatible]>;
124
125#endif // LLZK_SHARED_OP_HELPER