mlir/test/mlir-tblgen/rewriter-indexing.td

   1 // RUN: mlir-tblgen -gen-rewriters -I %S/../../include %s | FileCheck %s
   2
   3 include "mlir/IR/OpBase.td"
   4 include "mlir/IR/PatternBase.td"
   5
   6 def Test_Dialect : Dialect {
   7   let name = "test";
   8 }
   9 class NS_Op<string mnemonic, list<Trait> traits> :
  10     Op<Test_Dialect, mnemonic, traits>;
  11
  12 def AOp : NS_Op<"a_op", []> {
  13   let arguments = (ins
  14     AnyInteger:$any_integer
  15   );
  16
  17   let results = (outs AnyInteger);
  18 }
  19
  20 def BOp : NS_Op<"b_op", []> {
  21   let arguments = (ins
  22     AnyAttr: $any_attr,
  23     AnyInteger
  24   );
  25 }
  26
  27 // Tests dag operand indexing for ops with mixed attr and operand.
  28 // ---
  29
  30 def COp : NS_Op<"c_op", []> {
  31   let arguments = (ins
  32     AnyAttr: $any_attr1,
  33     AnyInteger,
  34     AnyAttr: $any_attr2,
  35     AnyInteger
  36   );
  37 }
  38
  39 // Only operand 0 should be addressed during matching.
  40 // CHECK: struct test1 : public ::mlir::RewritePattern {
  41 // CHECK: castedOp0.getODSOperands(0).begin()).getDefiningOp()
  42 def test1 : Pat<(BOp $attr, (AOp $input)),
  43                 (BOp $attr, $input)>;
  44
  45 // Only operand 0 and 1 should be addressed during matching.
  46 // CHECK: struct test2 : public ::mlir::RewritePattern {
  47 // CHECK: castedOp0.getODSOperands(0);
  48 // CHECK: castedOp0.getODSOperands(1).begin()).getDefiningOp()
  49 def test2 : Pat<(COp $attr1, $op1, $attr2, (AOp $op2)),
  50                 (BOp $attr1, $op2)>;
  51
  52
  53 // Check rewriting with a DAG subtree in the result and remapping a location.
  54 // CHECK: struct test3 : public ::mlir::RewritePattern {
  55 // We expect ODSOperand 0 here, the attribute before the operand in BOp
  56 // definition shouldn't shift the counter.
  57 // CHECK: op1 = (*castedOp0.getODSOperands(0).begin()).getDefiningOp();
  58 // CHECK: rewriter.create<test::BOp>((*a.getODSResults(0).begin()).getLoc()
  59 def test3 : Pat<(BOp $attr, (AOp:$a $input)),
  60                 (BOp $attr, (AOp $input), (location $a))>;
  61
  62 def DOp : NS_Op<"d_op", []> {
  63   let arguments = (ins
  64     AnyInteger:$v1,
  65     AnyInteger:$v2,
  66     AnyInteger:$v3,
  67     AnyInteger:$v4,
  68     AnyInteger:$v5,
  69     AnyInteger:$v6,
  70     AnyInteger:$v7,
  71     AnyInteger:$v8,
  72     AnyInteger:$v9,
  73     AnyInteger:$v10
  74   );
  75
  76   let results = (outs AnyInteger);
  77 }
  78
  79 def NativeBuilder :
  80   NativeCodeCall<[{
  81     nativeCall($_builder, $_loc, $0, $1, $2, $3, $4, $5, $6, $7, $8, $9)
  82   }]>;
  83
  84 // Check Pattern with large number of DAG arguments passed to NativeCodeCall
  85 // CHECK: struct test4 : public ::mlir::RewritePattern {
  86 // CHECK: nativeCall(rewriter, odsLoc, (*v1.begin()), (*v2.begin()), (*v3.begin()), (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
  87 def test4 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
  88                 (NativeBuilder $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;
  89
  90 // CHECK: struct test5 : public ::mlir::RewritePattern {
  91 // CHECK: foo(rewriter, (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
  92 def test5 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
  93                 (NativeCodeCall<[{ foo($_builder, $3...) }]> $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;
  94
  95 // Check Pattern with return type builder.
  96 def SameTypeAs : NativeCodeCall<"$0.getType()">;
  97 // CHECK: struct test6 : public ::mlir::RewritePattern {
  98 // CHECK: tblgen_types.push_back((*v2.begin()).getType())
  99 // CHECK: tblgen_types.push_back(rewriter.getI32Type())
 100 // CHECK: nativeVar_1 = doSomething((*v3.begin()))
 101 // CHECK: tblgen_types.push_back(nativeVar_1)
 102 def test6 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
 103                 (AOp (AOp $v1, (returnType $v2, "$_builder.getI32Type()", (NativeCodeCall<"doSomething($0)"> $v3))))>;