[Mlir-commits] [mlir] 853a614 - [mlir:OpFormatGen] Add Support for `$_ctxt` in the transformer.

[Sean Silva]

Author: Sean Silva
Date: 2021-06-14T18:02:55-07:00
New Revision: 853a614864754cd4b000f03a7ab8fbba103d6177

URL: https://github.com/llvm/llvm-project/commit/853a614864754cd4b000f03a7ab8fbba103d6177
DIFF: https://github.com/llvm/llvm-project/commit/853a614864754cd4b000f03a7ab8fbba103d6177.diff

LOG: [mlir:OpFormatGen] Add Support for `$_ctxt` in the transformer.

This is useful for "build tuple" type ops. In my case, in npcomp, I have
an op:

```
// Result type is `!torch.tuple<!torch.tensor, !torch.tensor>`.
torch.prim.TupleConstruct %0, %1 : !torch.tensor, !torch.tensor
```

and the context is required for the `Torch::TupleType::get` call (for
the case of an empty tuple).

The handling of these FmtContext's in the code is pretty ad-hoc -- I didn't
attempt to rationalize it and just made a targeted fix. As someone
unfamiliar with the code I had a hard time seeing how to more broadly fix
the situation.

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

Added: 
    

Modified: 
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/op-format.mlir
    mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
    mlir/tools/mlir-tblgen/OpFormatGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 847436ea3fb09..ea39b9c8c65ab 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1876,6 +1876,15 @@ def FormatTypesMatchAttrOp : TEST_Op<"format_types_match_attr", [
   let assemblyFormat = "attr-dict $value";
 }
 
+def FormatTypesMatchContextOp : TEST_Op<"format_types_match_context", [
+    TypesMatchWith<"tuple result type matches operand type", "value", "result",
+        "::mlir::TupleType::get($_ctxt, $_self)">
+  ]> {
+  let arguments = (ins AnyType:$value);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "attr-dict $value `:` type($value)";
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index e6f998fa4ac39..759e7f5e3abdd 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -348,3 +348,6 @@ test.format_infer_variadic_type_from_non_variadic %i64, %i64 : i64
 
 // CHECK: test.format_types_match_attr 1 : i64
 %ignored_res5 = test.format_types_match_attr 1 : i64
+
+// CHECK: test.format_types_match_context %[[I64]] : i64
+%ignored_res6 = test.format_types_match_context %i64 : i64

diff  --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 78e84d7dba9f0..cded9afe9dfdb 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -579,6 +579,7 @@ OpEmitter::OpEmitter(const Operator &op,
       opClass(op.getCppClassName(), op.getExtraClassDeclaration()),
       staticVerifierEmitter(staticVerifierEmitter) {
   verifyCtx.withOp("(*this->getOperation())");
+  verifyCtx.addSubst("_ctxt", "this->getOperation()->getContext()");
 
   genTraits();
 

diff  --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 3c3f00f2379c8..2c91708af5be5 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -1343,6 +1343,7 @@ void OperationFormat::genParserTypeResolution(Operator &op,
     } else if (const NamedTypeConstraint *var = resolver.getVariable()) {
       if (Optional<StringRef> tform = resolver.getVarTransformer()) {
         FmtContext fmtContext;
+        fmtContext.addSubst("_ctxt", "parser.getBuilder().getContext()");
         if (var->isVariadic())
           fmtContext.withSelf(var->name + "Types");
         else