[mlir][TestDialect] Fix invalid custom op printers

[rkayaith]

Summary

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

mlir/test/IR/parser.mlir

 // 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)>
 
 // CHECK-DAG: #map{{[0-9]*}} = affine_map<(d0) -> (d0)>
 #map1 = affine_map<(d0) -> (d0)>
 
 // CHECK-DAG: #map{{[0-9]*}} = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @special_float_values_in_tensors() {
 
   // CHECK: sparse<[{{\[}}1, 1, 0], [0, 1, 1]], [0xFFFFFFFF, 0x7F800001]>
   "foo"(){bar = sparse<[[1,1,0],[0,1,1]], [0xFFFFFFFF, 0x7F800001]> : tensor<2x2x2xf32>} : () -> ()
 }
 
 // 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"() : () -> ()
   }
   return
 }
 
 // Test allowing different name scopes for regions isolated from above.
 

 
 // CHECK-LABEL: func @parse_wrapped_keyword_test
 func.func @parse_wrapped_keyword_test() {
   // CHECK: test.parse_wrapped_keyword foo.keyword
   test.parse_wrapped_keyword foo.keyword
   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
 }
 
 // CHECK-LABEL: func @"\22_string_symbol_reference\22"
 func.func @"\"_string_symbol_reference\""() {
   // CHECK: ref = @"\22_string_symbol_reference\22"
   "foo.symbol_reference"() {ref = @"\"_string_symbol_reference\""} : () -> ()

mlir/test/IR/pretty-region-args.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
 
 // CHECK-LABEL: func @custom_region_names
 func.func @custom_region_names() -> () {
   "test.polyfor"() ({
   ^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
 }
 
 // CHECK-LABEL: func @weird_names
 // Make sure the asmprinter handles weird names correctly.
 func.func @weird_names() -> () {
   "test.polyfor"() ({
   ^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
 }
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/Base64.h"
 
 #include <cstdint>
 #include <numeric>
 #include <optional>
 
 // Include this before the using namespace lines below to
 // test that we don't have namespace dependencies.
 #include "TestOpsDialect.cpp.inc"
if (parser.parseOperand(argInfo.ssaName) ||
     return failure();
 
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
   return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);
 }
 
 void IsolatedRegionOp::print(OpAsmPrinter &p) {
-  p << "test.isolated_region ";
+  p << ' ';
   p.printOperand(getOperand());
   p.shadowRegionArgs(getRegion(), getOperand());
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Test SSACFGRegionOp

 
 ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
   return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
 }
 
 void AffineScopeOp::print(OpAsmPrinter &p) {
-  p << "test.affine_scope ";
+  p << " ";
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Test parser.
 //===----------------------------------------------------------------------===//
 
 ParseResult ParseIntegerLiteralOp::parse(OpAsmParser &parser,
ParseResult ParseB64BytesOp::parse(OpAsmParser &parser,
   if (parser.parseBase64Bytes(&bytes))
     return failure();
   result.addAttribute("b64", parser.getBuilder().getStringAttr(
                                  StringRef(&bytes.front(), bytes.size())));
   return success();
 }
 
 void ParseB64BytesOp::print(OpAsmPrinter &p) {
-  // Don't print the base64 version to check that we decoded it correctly.
-  p << " \"" << getB64() << "\"";
+  p << " \"" << llvm::encodeBase64(getB64()) << "\"";
 }
 
 //===----------------------------------------------------------------------===//
 // Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
 
 ParseResult WrappingRegionOp::parse(OpAsmParser &parser,
                                     OperationState &result) {
   if (parser.parseKeyword("wraps"))
ParseResult PolyForOp::parse(OpAsmParser &parser, OperationState &result) {
 
   // Parse the body region.
   Region *body = result.addRegion();
   for (auto &iv : ivsInfo)
     iv.type = parser.getBuilder().getIndexType();
   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) {
   auto arrayAttr = getOperation()->getAttrOfType<ArrayAttr>("arg_names");
   if (!arrayAttr)
     return;
   auto args = getRegion().front().getArguments();
   auto e = std::min(arrayAttr.size(), args.size());