mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

   1 # RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS
   2 # RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL
   3
   4 # @linalg_structured_op
   5 # def test1(O=TensorDef(T, S.M, S.N, output=True),
   6 #           cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
   7 #   """Title.
   8
   9 #   Detailed description.
  10 #   """
  11 #   O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n))
  12
  13 --- !LinalgOpConfig
  14 metadata: !LinalgOpMetadata
  15   name: test1
  16   cpp_class_name: Test1Op
  17   doc: |-
  18     Title.
  19
  20     Detailed description.
  21 structured_op: !LinalgStructuredOpConfig
  22   args:
  23   - !LinalgOperandDefConfig
  24     name: O
  25     kind: output_tensor
  26     type_var: T
  27     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
  28   - !LinalgOperandDefConfig
  29     name: cast
  30     kind: type_fn_attr
  31     default_fn: cast_signed
  32   indexing_maps: !LinalgIndexingMapsConfig
  33     static_indexing_maps:
  34     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
  35   iterator_types:
  36   - parallel
  37   - parallel
  38   assignments:
  39   - !ScalarAssign
  40     arg: O
  41     value: !ScalarExpression
  42       scalar_fn:
  43         kind: binary
  44         fn_name: add
  45         operands:
  46         - !ScalarExpression
  47           scalar_fn:
  48             kind: type
  49             attr_name: cast
  50             type_var: T
  51             operands:
  52             - !ScalarExpression
  53               scalar_const: '42 : i64'
  54         - !ScalarExpression
  55           scalar_fn:
  56             kind: type
  57             attr_name: cast
  58             type_var: T
  59             operands:
  60             - !ScalarExpression
  61               scalar_index: 1
  62
  63 # ODS-LABEL:  def Test1Op : LinalgStructuredBase_Op<"test1"
  64
  65 #       ODS:  let summary = [{Title.}];
  66 #  ODS-NEXT:  let description = [{Detailed description.}];
  67
  68 #       ODS:  let arguments =
  69 #  ODS-NEXT:    Variadic<AnyType>:$inputs,
  70 #  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
  71 #  ODS-NEXT:    DefaultValuedOptionalAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
  72
  73 #       ODS:  let builders =
  74 #       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
  75 #  ODS-NEXT:       "ValueRange":$outputs,
  76 #  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
  77
  78 #       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
  79 #  ODS-NEXT:       "ValueRange":$outputs, "Attribute":$cast,
  80 #  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
  81
  82 #       ODS:    buildStructuredOp($_builder, $_state, resultTensorTypes,
  83 #  ODS-NEXT:      attributes, Test1Op::getRegionBuilder())
  84
  85 #       ODS:    MutableOperandRange getDpsInitsMutable() {
  86 #  ODS-NEXT:      return getOutputsMutable()
  87 #  ODS-NEXT:    }
  88
  89 # IMPL-LABEL:  void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
  90 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
  91 #       IMPL:  TypeFn castVal = TypeFn::cast_signed;
  92 #  IMPL-NEXT:  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
  93 #  IMPL-NEXT:                                return attr.getName() == "cast"; });
  94 #  IMPL-NEXT:  if (castIter != attrs.end()) {
  95 #  IMPL-NEXT:    if (auto attr = llvm::dyn_cast<TypeFnAttr>(castIter->getValue()))
  96 #  IMPL-NEXT:      castVal = attr.getValue();
  97 #  IMPL-NEXT:  }
  98
  99 #       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
 100 #   IMPL-DAG:  Value [[VAL1:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL0]]);
 101 #   IMPL-DAG:  Value [[VAL2:[a-z0-9]+]] = helper.index(1);
 102 #   IMPL-DAG:  Value [[VAL3:[a-z0-9]+]] = helper.buildTypeFn(castVal, block.getArgument(0).getType(), [[VAL2]]);
 103 #   IMPL-DAG:  Value [[VAL4:[a-z0-9]+]] = helper.buildBinaryFn(BinaryFn::add, [[VAL1]], [[VAL3]]);
 104
 105
 106 # @linalg_structured_op
 107 # def test2(I=TensorDef(T, S.M, S.N),
 108 #           O=TensorDef(T, S.M, S.N, output=True),
 109 #           strides=IndexAttrDef(S.SM, S.SN, default=[1, 2])):
 110 #   """Title.
 111
 112 #   Detailed description.
 113 #   """
 114 #   O[D.m, D.n] = I[D.n * S.SM, D.m * S.SN]
 115
 116 --- !LinalgOpConfig
 117 metadata: !LinalgOpMetadata
 118   name: test2
 119   cpp_class_name: Test2Op
 120   doc: |-
 121     Title.
 122
 123     Detailed description.
 124 structured_op: !LinalgStructuredOpConfig
 125   args:
 126   - !LinalgOperandDefConfig
 127     name: I
 128     kind: input_tensor
 129     type_var: T
 130     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
 131   - !LinalgOperandDefConfig
 132     name: O
 133     kind: output_tensor
 134     type_var: T
 135     shape_map: affine_map<()[s0, s1, s2, s3] -> (s0, s1)>
 136   - !LinalgOperandDefConfig
 137     name: strides
 138     kind: index_attr
 139     index_attr_map: affine_map<()[s0, s1, s2, s3] -> (s2, s3)>
 140     default_indices:
 141     - 1
 142     - 2
 143   indexing_maps: !LinalgIndexingMapsConfig
 144     static_indexing_maps:
 145     - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>
 146     - affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>
 147   iterator_types:
 148   - parallel
 149   - parallel
 150   assignments:
 151   - !ScalarAssign
 152     arg: O
 153     value: !ScalarExpression
 154       scalar_arg: I
 155
 156 # ODS-LABEL:  def Test2Op : LinalgStructuredBase_Op<"test2"
 157
 158 #       ODS:  let arguments =
 159 #  ODS-NEXT:    Variadic<AnyType>:$inputs,
 160 #  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
 161 #  ODS-NEXT:    DefaultValuedOptionalAttr<RankedI64ElementsAttr<[2]>
 162 #  ODS-SAME:    "{ static_cast<int64_t>(1), static_cast<int64_t>(2) }">:$strides
 163
 164 #       ODS:  "Attribute":$strides
 165 #       ODS:  $_state.addAttribute("strides", strides);
 166
 167 #       ODS:  bool hasDynamicIndexingMaps();
 168 #  ODS-NEXT:  LogicalResult verifyIndexingMapRequiredAttributes();
 169
 170 #       IMPL:  getSymbolBindings(Test2Op self)
 171 #       IMPL:  cst2 = self.getStrides().getValues<int64_t>()[0];
 172 #  IMPL-NEXT:  getAffineConstantExpr(cst2, context)
 173 #       IMPL:  cst3 = self.getStrides().getValues<int64_t>()[1];
 174 #  IMPL-NEXT:  getAffineConstantExpr(cst3, context)
 175
 176 #       IMPL:  Test2Op::getIndexingMaps()
 177 #       IMPL:  = getSymbolBindings(*this);
 178 #       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d1 * s2, d0 * s3)>"
 179 #       IMPL:  "affine_map<(d0, d1)[s0, s1, s2, s3] -> (d0, d1)>"
 180
 181 #       IMPL:  Test2Op::getNumRegionArgs() { return 2; }
 182
 183 #       IMPL:  Test2Op::hasDynamicIndexingMaps() { return true; }
 184 #       IMPL:  Test2Op::verifyIndexingMapRequiredAttributes()
 185 #       IMPL:  auto attr = op->getAttrOfType<DenseElementsAttr>("strides")
 186 #       IMPL:  "incorrect element type for index attribute 'strides'"
 187 #       IMPL:  "incorrect shape for index attribute 'strides'"
 188 #       IMPL:  void Test2Op::regionBuilder(ImplicitLocOpBuilder &b,
 189 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
 190 #  IMPL-NEXT:    assert(2 > 0 && block.getNumArguments() == 2 &&
 191
 192 #       IMPL:   yields.push_back(block.getArgument(0));
 193
 194 # @linalg_structured_op
 195 # def test3(value=ScalarDef(T1),
 196 #           O=TensorDef(U, output=True)):
 197 #   """Title.
 198
 199 #   Detailed description.
 200 #   """
 201 #   O[None] = TypeFn.cast_signed(U, value)
 202
 203 --- !LinalgOpConfig
 204 metadata: !LinalgOpMetadata
 205   name: test3
 206   cpp_class_name: Test3Op
 207   doc: |-
 208     Title.
 209
 210     Detailed description.
 211 structured_op: !LinalgStructuredOpConfig
 212   args:
 213   - !LinalgOperandDefConfig
 214     name: value
 215     kind: scalar
 216     type_var: T1
 217   - !LinalgOperandDefConfig
 218     name: O
 219     kind: output_tensor
 220     type_var: U
 221     shape_map: affine_map<() -> ()>
 222   indexing_maps: !LinalgIndexingMapsConfig
 223     static_indexing_maps:
 224     - affine_map<() -> ()>
 225     - affine_map<() -> ()>
 226   iterator_types: []
 227   assignments:
 228   - !ScalarAssign
 229     arg: O
 230     value: !ScalarExpression
 231       scalar_fn:
 232         kind: type
 233         fn_name: cast_signed
 234         type_var: U
 235         operands:
 236         - !ScalarExpression
 237           scalar_arg: value
 238
 239 #       IMPL:  Test3Op::getIteratorTypesArray() {
 240 #  IMPL-NEXT:    int64_t rank = getRank(getDpsInitOperand(0));
 241
 242 #       IMPL:  Test3Op::getIndexingMaps() {
 243 #  IMPL-NEXT:    MLIRContext *context = getContext();
 244 #  IMPL-NEXT:    AffineMap scalarMap = AffineMap::get(getNumParallelLoops(), 0, context);
 245 #  IMPL-NEXT:    AffineMap tensorMap = AffineMap::getMultiDimIdentityMap(
 246
 247 # @linalg_structured_op
 248 # def test4(O=TensorDef(T, S.M, S.N, output=True),
 249 #           unary_fun=UnaryFnAttrDef(default=UnaryFn.exp),
 250 #           binary_fun=BinaryFnAttrDef(default=BinaryFn.add)):
 251 #   """Title.
 252
 253 #   Detailed description.
 254 #   """
 255 #   O[D.m, D.n] = binary_fun(unary_fun(O[D.m, D.n]), O[D.m, D.n])
 256
 257 --- !LinalgOpConfig
 258 metadata: !LinalgOpMetadata
 259   name: test4
 260   cpp_class_name: Test4Op
 261   doc: |-
 262     Title.
 263
 264     Detailed description.
 265 structured_op: !LinalgStructuredOpConfig
 266   args:
 267   - !LinalgOperandDefConfig
 268     name: O
 269     kind: output_tensor
 270     type_var: T
 271     shape_map: affine_map<()[s0, s1] -> (s0, s1)>
 272   - !LinalgOperandDefConfig
 273     name: unary_fun
 274     kind: unary_fn_attr
 275     default_fn: exp
 276   - !LinalgOperandDefConfig
 277     name: binary_fun
 278     kind: binary_fn_attr
 279     default_fn: add
 280   indexing_maps: !LinalgIndexingMapsConfig
 281     static_indexing_maps:
 282     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
 283   iterator_types:
 284   - parallel
 285   - parallel
 286   assignments:
 287   - !ScalarAssign
 288     arg: O
 289     value: !ScalarExpression
 290       scalar_fn:
 291         kind: binary
 292         attr_name: binary_fun
 293         operands:
 294         - !ScalarExpression
 295           scalar_fn:
 296             kind: unary
 297             attr_name: unary_fun
 298             operands:
 299             - !ScalarExpression
 300               scalar_arg: O
 301         - !ScalarExpression
 302           scalar_arg: O
 303
 304 # ODS-LABEL:  def Test4Op : LinalgStructuredBase_Op<"test4"
 305
 306 #       ODS:  let arguments =
 307 #  ODS-NEXT:    Variadic<AnyType>:$inputs,
 308 #  ODS-NEXT:    Variadic<AnyShaped>:$outputs,
 309 #  ODS-NEXT:    DefaultValuedOptionalAttr<UnaryFnAttr, "UnaryFn::exp">:$unary_fun,
 310 #  ODS-NEXT:    DefaultValuedOptionalAttr<BinaryFnAttr, "BinaryFn::add">:$binary_fun
 311
 312 #       ODS:    "Attribute":$unary_fun, "Attribute":$binary_fun,
 313
 314 #       ODS:    $_state.addAttribute("unary_fun", unary_fun)
 315 #  ODS-NEXT:    $_state.addAttribute("binary_fun", binary_fun)
 316
 317 # IMPL-LABEL:  void Test4Op::regionBuilder(ImplicitLocOpBuilder &b,
 318 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
 319 #       IMPL:  UnaryFn unary_funVal = UnaryFn::exp
 320 #       IMPL:  BinaryFn binary_funVal = BinaryFn::add
 321
 322 #       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.buildUnaryFn(unary_funVal, block.getArgument(0))
 323 #  IMPL-NEXT:  Value [[VAL1:[a-z0-9]+]] = helper.buildBinaryFn(binary_funVal, [[VAL0]], block.getArgument(0))
 324 #  IMPL-NEXT:  yields.push_back([[VAL1]])
 325
 326 # @linalg_structured_op
 327 # def test5(value=ScalarDef(T1), O=TensorDef(U, output=True)):
 328 #   """Title.
 329
 330 #   Detailed description.
 331 #   """
 332 #   implements(FillOpInterface)
 333 #   defines(Canonicalizer)
 334 #   O[None] = TypeFn.cast(U, value)
 335
 336 --- !LinalgOpConfig
 337 metadata: !LinalgOpMetadata
 338   name: test5
 339   cpp_class_name: Test5Op
 340   doc: |-
 341     Title.
 342
 343     Detailed description.
 344   implements:
 345   - LinalgFillOpInterface
 346   defines:
 347   - hasCanonicalizer
 348 structured_op: !LinalgStructuredOpConfig
 349   args:
 350   - !LinalgOperandDefConfig
 351     name: value
 352     kind: scalar
 353     type_var: T1
 354   - !LinalgOperandDefConfig
 355     name: O
 356     kind: output_tensor
 357     type_var: U
 358     shape_map: affine_map<() -> ()>
 359   indexing_maps: !LinalgIndexingMapsConfig
 360     static_indexing_maps:
 361     - affine_map<() -> ()>
 362     - affine_map<() -> ()>
 363   iterator_types: []
 364   assignments:
 365   - !ScalarAssign
 366     arg: O
 367     value: !ScalarExpression
 368       scalar_fn:
 369         kind: type
 370         fn_name: cast
 371         type_var: U
 372         operands:
 373         - !ScalarExpression
 374           scalar_arg: value
 375
 376 # ODS-LABEL:  def Test5Op : LinalgStructuredBase_Op<"test5"
 377 #  ODS-NEXT:  /*extraInterfaces=*/[LinalgFillOpInterface])>
 378
 379 #       ODS:  let hasCanonicalizer = 1;