[Mlir-commits] [mlir] b89b3cd - [mlir][TestDialect] Fix invalid custom op printers

[Rahul Kayaith]

Author: Rahul Kayaith
Date: 2023-07-04T23:29:40-04:00
New Revision: b89b3cdb612f43e91c7330a13c3f60a1ad251e50

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

LOG: [mlir][TestDialect] Fix invalid custom op printers

This fixes a few custom printers which were printing IR that couldn't be
round-tripped.

Reviewed By: rriddle

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

Added: 
    

Modified: 
    mlir/test/IR/parser.mlir
    mlir/test/IR/pretty-region-args.mlir
    mlir/test/lib/Dialect/Test/TestDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index 2041ca90e8116e..66c9adca8f98ce 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -1,4 +1,7 @@
 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | FileCheck %s -check-prefix GENERIC
+// RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -mlir-print-op-generic | FileCheck %s -check-prefix GENERIC
 
 // CHECK-DAG: #map{{[0-9]*}} = affine_map<(d0, d1, d2, d3, d4)[s0] -> (d0, d1, d2, d4, d3)>
 #map = affine_map<(d0, d1, d2, d3, d4)[s0] -> (d0, d1, d2, d4, d3)>
@@ -1128,10 +1131,10 @@ func.func @special_float_values_in_tensors() {
 // Test parsing of an op with multiple region arguments, and without a
 // delimiter.
 
-// CHECK-LABEL: func @op_with_region_args
+// GENERIC-LABEL: op_with_region_args
 func.func @op_with_region_args() {
-  // CHECK: "test.polyfor"() ({
-  // CHECK-NEXT: ^bb{{.*}}(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index):
+  // GENERIC: "test.polyfor"() ({
+  // GENERIC-NEXT: ^bb{{.*}}(%{{.*}}: index, %{{.*}}: index, %{{.*}}: index):
   test.polyfor %i, %j, %k {
     "foo"() : () -> ()
   }
@@ -1185,9 +1188,9 @@ func.func @parse_wrapped_keyword_test() {
   return
 }
 
-// CHECK-LABEL: func @parse_base64_test
+// GENERIC-LABEL: parse_base64_test
 func.func @parse_base64_test() {
-  // CHECK: test.parse_b64 "hello world"
+  // GENERIC: "test.parse_b64"() <{b64 = "hello world"}>
   test.parse_b64 "aGVsbG8gd29ybGQ="
   return
 }

diff  --git a/mlir/test/IR/pretty-region-args.mlir b/mlir/test/IR/pretty-region-args.mlir
index 764bea2ee0d208..f6420142b2372d 100644
--- a/mlir/test/IR/pretty-region-args.mlir
+++ b/mlir/test/IR/pretty-region-args.mlir
@@ -6,8 +6,7 @@ func.func @custom_region_names() -> () {
   ^bb0(%arg0: index, %arg1: index, %arg2: index):
     "foo"() : () -> ()
   }) { arg_names = ["i", "j", "k"] } : () -> ()
-  // CHECK: test.polyfor
-  // CHECK-NEXT: ^bb{{.*}}(%i: index, %j: index, %k: index):
+  // CHECK: test.polyfor %i, %j, %k
   return
 }
 
@@ -18,8 +17,7 @@ func.func @weird_names() -> () {
   ^bb0(%arg0: i32, %arg1: i32, %arg2: index):
     "foo"() : () -> i32
   }) { arg_names = ["a .^x", "0"] } : () -> ()
-  // CHECK: test.polyfor
-  // CHECK-NEXT: ^bb{{.*}}(%a_.5Ex: i32, %_0: i32, %arg0: index):
+  // CHECK: test.polyfor %a_.5Ex, %_0, %arg0
   // CHECK-NEXT: %0 = "foo"()
   return
 }

diff  --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index d167d0274aad68..1ec769847a1c44 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -34,6 +34,7 @@
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/Base64.h"
 
 #include <cstdint>
 #include <numeric>
@@ -1020,7 +1021,7 @@ ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
 }
 
 void IsolatedRegionOp::print(OpAsmPrinter &p) {
-  p << "test.isolated_region ";
+  p << ' ';
   p.printOperand(getOperand());
   p.shadowRegionArgs(getRegion(), getOperand());
   p << ' ';
@@ -1054,7 +1055,7 @@ ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
 }
 
 void AffineScopeOp::print(OpAsmPrinter &p) {
-  p << "test.affine_scope ";
+  p << " ";
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
 }
 
@@ -1103,8 +1104,7 @@ ParseResult ParseB64BytesOp::parse(OpAsmParser &parser,
 }
 
 void ParseB64BytesOp::print(OpAsmPrinter &p) {
-  // Don't print the base64 version to check that we decoded it correctly.
-  p << " \"" << getB64() << "\"";
+  p << " \"" << llvm::encodeBase64(getB64()) << "\"";
 }
 
 //===----------------------------------------------------------------------===//
@@ -1260,7 +1260,14 @@ ParseResult PolyForOp::parse(OpAsmParser &parser, OperationState &result) {
   return parser.parseRegion(*body, ivsInfo);
 }
 
-void PolyForOp::print(OpAsmPrinter &p) { p.printGenericOp(*this); }
+void PolyForOp::print(OpAsmPrinter &p) {
+  p << " ";
+  llvm::interleaveComma(getRegion().getArguments(), p, [&](auto arg) {
+    p.printRegionArgument(arg, /*argAttrs =*/{}, /*omitType=*/true);
+  });
+  p << " ";
+  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
+}
 
 void PolyForOp::getAsmBlockArgumentNames(Region &region,
                                          OpAsmSetValueNameFn setNameFn) {