[Mlir-commits] [mlir] 0164119 - [MLIR] Remove TableGen redundant calls to native calls when creating new operations in DRR TableGen files

[Jacques Pienaar]

Author: AlexEichenberger
Date: 2020-06-22T08:12:04-07:00
New Revision: 01641197ee0bd3895c756c925723c8c8e03bcb09

URL: https://github.com/llvm/llvm-project/commit/01641197ee0bd3895c756c925723c8c8e03bcb09
DIFF: https://github.com/llvm/llvm-project/commit/01641197ee0bd3895c756c925723c8c8e03bcb09.diff

LOG: [MLIR] Remove TableGen redundant calls to native calls when creating new operations in DRR TableGen files

Summary:
Currently, the TableGen rewrite generates redundant native calls in MLIR DRR files. This is a problem as some native calls may involve significant computations (e.g. when performing constant propagation where every values in a large tensor is touched).

The pattern was as follow:

```c++
if (native-call(args)) tblgen_attrs.emplace_back(rewriter, attribute, native-call(args))
```

The replacement pattern compute `native-call(args)` once and then use it both in the `if` condition and the `emplace_back` call.

Differential Revision: https://reviews.llvm.org/D82101

Added: 
    

Modified: 
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/lib/Dialect/Test/TestPatterns.cpp
    mlir/test/mlir-tblgen/pattern.mlir
    mlir/tools/mlir-tblgen/RewriterGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 630a7c58e449..6aa18c1324bd 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -602,6 +602,20 @@ def OpJ : TEST_Op<"op_j">, Arguments<(ins)>, Results<(outs I32)>;
 def OpK : TEST_Op<"op_k">, Arguments<(ins)>, Results<(outs I32)>;
 def : Pat<(OpJ), (OpK)>;
 
+// Test that natives calls are only called once during rewrites.
+def OpM : TEST_Op<"op_m"> {
+  let arguments = (ins I32, OptionalAttr<I32Attr>:$optional_attr);
+  let results = (outs I32);
+}
+// Pattern add the argument plus a increasing static number hidden in
+// OpMTest function. That value is set into the optional argument.
+// That way, we will know if operations is called once or twice.
+def OpMGetNullAttr : NativeCodeCall<"Attribute()">;
+def OpMAttributeIsNull : Constraint<CPred<"! ($_self)">, "Attribute is null">;
+def OpMVal : NativeCodeCall<"OpMTest($_builder, $0)">;
+def : Pat<(OpM $attr, $optAttr), (OpM $attr, (OpMVal $attr) ),
+    [(OpMAttributeIsNull:$optAttr)]>;
+
 // Test `$_` for ignoring op argument match.
 def TestIgnoreArgMatchSrcOp : TEST_Op<"ignore_arg_match_src"> {
   let arguments = (ins

diff  --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 60f663f75adc..f44b987f17cb 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -32,6 +32,16 @@ static void handleNoResultOp(PatternRewriter &rewriter,
                                     op.operand());
 }
 
+// Test that natives calls are only called once during rewrites.
+// OpM_Test will return Pi, increased by 1 for each subsequent calls.
+// This let us check the number of times OpM_Test was called by inspecting
+// the returned value in the MLIR output.
+static int64_t opMIncreasingValue = 314159265;
+static Attribute OpMTest(PatternRewriter &rewriter, Value val) {
+  int64_t i = opMIncreasingValue++;
+  return rewriter.getIntegerAttr(rewriter.getIntegerType(32), i);
+}
+
 namespace {
 #include "TestPatterns.inc"
 } // end anonymous namespace

diff  --git a/mlir/test/mlir-tblgen/pattern.mlir b/mlir/test/mlir-tblgen/pattern.mlir
index 6154e6bc4c45..0f2f434c928f 100644
--- a/mlir/test/mlir-tblgen/pattern.mlir
+++ b/mlir/test/mlir-tblgen/pattern.mlir
@@ -359,3 +359,14 @@ func @generateVariadicOutputOpInNestedPattern() -> (i32) {
   %0 = "test.one_i32_out"() : () -> (i32)
   return %0 : i32
 }
+
+//===----------------------------------------------------------------------===//
+// Test that natives calls are only called once during rewrites.
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: redundantTest
+func @redundantTest(%arg0: i32) -> i32 {
+  %0 = "test.op_m"(%arg0) : (i32) -> i32
+  // CHECK: "test.op_m"(%arg0) {optional_attr = 314159265 : i32} : (i32) -> i32
+  return %0 : i32
+}

diff  --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 62638e6c9964..13498e06e54e 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1044,11 +1044,11 @@ void PatternEmitter::createAggregateLocalVarsForOpArgs(
   os.indent(6) << formatv(
       "SmallVector<NamedAttribute, 4> tblgen_attrs; (void)tblgen_attrs;\n");
 
+  const char *addAttrCmd =
+      "if (auto tmpAttr = {1}) "
+      "tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), tmpAttr);\n";
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
-      const char *addAttrCmd = "if ({1}) {{"
-                               "  tblgen_attrs.emplace_back(rewriter."
-                               "getIdentifier(\"{0}\"), {1}); }\n";
       // The argument in the op definition.
       auto opArgName = resultOp.getArgName(argIndex);
       if (auto subTree = node.getArgAsNestedDag(argIndex)) {