3 files changed, 59 insertions, 1 deletions
diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index dc85fbb14b3..a6460b46dd9 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -1116,3 +1116,18 @@ func @"\"_string_symbol_reference\""() {
 // CHECK-LABEL: func @nested_reference
 // CHECK-NEXT: ref = @some_symbol::@some_nested_symbol
 func @nested_reference() attributes {test.ref = @some_symbol::@some_nested_symbol }
+
+// CHECK-LABEL: func @custom_asm_names
+func @custom_asm_names() -> (i32, i32, i32, i32, i32, i32, i32) {
+  // CHECK: %[[FIRST:first.*]], %[[MIDDLE:middle_results.*]]:2, %[[LAST:[0-9]+]]
+  %0, %1:2, %2 = "test.asm_interface_op"() : () -> (i32, i32, i32, i32)
+
+  // CHECK: %[[FIRST_2:first.*]], %[[LAST_2:[0-9]+]]
+  %3, %4 = "test.asm_interface_op"() : () -> (i32, i32)
+
+  // CHECK: %[[RESULT:result.*]]
+  %5 = "test.asm_dialect_interface_op"() : () -> (i32)
+
+  // CHECK: return %[[FIRST]], %[[MIDDLE]]#0, %[[MIDDLE]]#1, %[[LAST]], %[[FIRST_2]], %[[LAST_2]]
+  return %0, %1#0, %1#1, %2, %3, %4, %5 : i32, i32, i32, i32, i32, i32, i32
+}
diff --git a/mlir/test/lib/TestDialect/TestDialect.cpp b/mlir/test/lib/TestDialect/TestDialect.cpp
index 01780432a1a..d838f75f7e7 100644
--- a/mlir/test/lib/TestDialect/TestDialect.cpp
+++ b/mlir/test/lib/TestDialect/TestDialect.cpp
@@ -30,6 +30,18 @@ using namespace mlir;
 //===----------------------------------------------------------------------===//
 
 namespace {
+
+// Test support for interacting with the AsmPrinter.
+struct TestOpAsmInterface : public OpAsmDialectInterface {
+  using OpAsmDialectInterface::OpAsmDialectInterface;
+
+  void getAsmResultNames(Operation *op,
+                         OpAsmSetValueNameFn setNameFn) const final {
+    if (auto asmOp = dyn_cast<AsmDialectInterfaceOp>(op))
+      setNameFn(asmOp, "result");
+  }
+};
+
 struct TestOpFolderDialectInterface : public OpFolderDialectInterface {
   using OpFolderDialectInterface::OpFolderDialectInterface;
 
@@ -112,7 +124,8 @@ TestDialect::TestDialect(MLIRContext *context)
 #define GET_OP_LIST
 #include "TestOps.cpp.inc"
       >();
-  addInterfaces<TestOpFolderDialectInterface, TestInlinerInterface>();
+  addInterfaces<TestOpAsmInterface, TestOpFolderDialectInterface,
+                TestInlinerInterface>();
   allowUnknownOperations();
 }
 
@@ -227,6 +240,7 @@ static void print(OpAsmPrinter &p, WrappingRegionOp op) {
 //===----------------------------------------------------------------------===//
 // Test PolyForOp - parse list of region arguments.
 //===----------------------------------------------------------------------===//
+
 static ParseResult parsePolyForOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 4> ivsInfo;
   // Parse list of region arguments without a delimiter.
@@ -241,6 +255,21 @@ static ParseResult parsePolyForOp(OpAsmParser &parser, OperationState &result) {
 }
 
 //===----------------------------------------------------------------------===//
+// Test OpAsmInterface.
+//===----------------------------------------------------------------------===//
+
+void AsmInterfaceOp::getAsmResultNames(
+    function_ref<void(Value *, StringRef)> setNameFn) {
+  // Give a name to the first and middle results.
+  setNameFn(firstResult(), "first");
+  if (!llvm::empty(middleResults()))
+    setNameFn(*middleResults().begin(), "middle_results");
+
+  // Use default numbering for the last result.
+  setNameFn(getResult(getNumResults() - 1), "");
+}
+
+//===----------------------------------------------------------------------===//
 // Test removing op with inner ops.
 //===----------------------------------------------------------------------===//
 
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index a0e1cd61ba4..1ccbda2f9b1 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -19,6 +19,7 @@
 #define TEST_OPS
 
 include "mlir/IR/OpBase.td"
+include "mlir/IR/OpAsmInterface.td"
 include "mlir/Analysis/CallInterfaces.td"
 include "mlir/Analysis/InferTypeOpInterface.td"
 
@@ -995,4 +996,17 @@ def PolyForOp : TEST_Op<"polyfor">
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test OpAsmInterface.
+
+def AsmInterfaceOp : TEST_Op<"asm_interface_op",
+    [DeclareOpInterfaceMethods<OpAsmOpInterface>]> {
+  let results = (outs AnyType:$firstResult, Variadic<AnyType>:$middleResults,
+                      AnyType);
+}
+
+def AsmDialectInterfaceOp : TEST_Op<"asm_dialect_interface_op"> {
+  let results = (outs AnyType);
+}
+
 #endif // TEST_OPS