Add support for two return values

This commit adds support for two return values of type
memref f32 and i64.

Signed-Off By: Vivek Khandelwal <[email protected]>

Author: vivekkhandelwal1

e2e_testing/torchscript/basic.py

@@ -998,3 +998,19 @@ def forward(self):
 def TensorOpaqueLiteralModule_basic(module, tu: TestUtils):
     module.forward()
 
+class ReturnTwoTensorF32I64(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+        ([-1, -1], torch.int64, True),
+    ])
+    def forward(self, a, b):
+        return a, b
+
+@register_test_case(module_factory=lambda: ReturnTwoTensorF32I64())
+def ReturnTwoTensorF32I64_basic(module, tu: TestUtils):
+    module.forward(tu.rand(2, 3), torch.randint(5, (2, 3)))

e2e_testing/torchscript/xfail_sets.py

@@ -49,4 +49,5 @@
     "SqueezeDimModule_static",
     "SqueezeDimModule_identity",
     "SqueezeDimModule_unitDim",
+    "ReturnTwoTensorF32I64_basic",
 }

lib/RefBackend/RefBackend.cpp

@@ -163,9 +163,9 @@ static LogicalResult mungeFunction(
     auto supportedFuncsEnd = supportedConsumeFuncReturnFuncs.end();
     std::string funcName = getConsumeReturnFunctionNameForReturnTypes(retTypes);
     if (supportedConsumeFuncReturnFuncs.find(funcName) == supportedFuncsEnd) {
-      op.emitError(
-          "must have one return value of memref types or scalar types "
-          "of i32, i64, f32, f64, i1, or three return values of memref f32");
+      op.emitError("must have one return value of memref types or scalar types "
+                   "of i32, i64, f32, f64, i1, or two return values of memref "
+                   "f32 and i64, or three return values of memref f32");
       isSupported = false;
     }
 
@@ -194,7 +194,8 @@ static std::set<std::string> getSupportedConsumeFuncReturnFuncs(OpBuilder &b) {
   Type f64 = b.getF64Type();
 
   SmallVector<TypeRange> supportedReturnTypes = {
-      mri1, mri32, mri64, mrf32, mrf64, i64, f32, f64, {mrf32, mrf32, mrf32}};
+      mri1, mri32, mri64, mrf32,          mrf64,
+      i64,  f32,   f64,   {mrf32, mri64}, {mrf32, mrf32, mrf32}};
 
   llvm::for_each(supportedReturnTypes, [&](TypeRange &types) {
     funcNames.insert(getConsumeReturnFunctionNameForReturnTypes(types));

python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py

@@ -65,6 +65,14 @@ def consume_return_f32(a):
         def consume_return_f64(a):
             self.result = a
 
+        @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor),
+                          ctypes.POINTER(UnrankedMemRefDescriptor))
+        def consume_return_mrf32_mri64(arg0, arg1):
+            self.result = unranked_memref_to_numpy(
+                arg0, np.float32), unranked_memref_to_numpy(
+                    arg1,
+                    np.int64)
+
         @ctypes.CFUNCTYPE(None, ctypes.POINTER(UnrankedMemRefDescriptor),
                           ctypes.POINTER(UnrankedMemRefDescriptor),
                           ctypes.POINTER(UnrankedMemRefDescriptor))
@@ -98,6 +106,10 @@ def consume_return_mrf32_mrf32_mrf32(arg0, arg1, arg2):
         self.ee.register_runtime("refbackend_consume_func_return_f64",
                                  consume_return_f64)
 
+        self.ee.register_runtime(
+            "refbackend_consume_func_return_mrf32_mri64",
+            consume_return_mrf32_mri64)
+
         self.ee.register_runtime(
             "refbackend_consume_func_return_mrf32_mrf32_mrf32",
             consume_return_mrf32_mrf32_mrf32)

test/RefBackend/munge-calling-conventions.mlir

@@ -53,3 +53,20 @@ func @elemental_type(%arg0: memref<i64>) -> i64 {
 func @multiple_return_values(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) -> (memref<?xf32>, memref<?xf32>, memref<?xf32>) {
   return %arg0 ,%arg1, %arg2 : memref<?xf32>, memref<?xf32>, memref<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL:   func @two_return_values(
+// CHECK-SAME:                                 %[[ARG0:.*]]: memref<*xf32>, %[[ARG1:.*]]: memref<*xi64>)
+// CHECK-SAME:                                 attributes {llvm.emit_c_interface} {
+// CHECK:           %[[VAL0:.*]] = memref.cast %[[ARG0]] : memref<*xf32> to memref<?xf32>
+// CHECK:           %[[VAL1:.*]] = memref.cast %[[ARG1]] : memref<*xi64> to memref<?xi64>
+// CHECK:           %[[RET0:.*]] = memref.cast %[[VAL0]] : memref<?xf32> to memref<*xf32>
+// CHECK:           %[[RET1:.*]] = memref.cast %[[VAL1]] : memref<?xi64> to memref<*xi64>
+// CHECK:           call @refbackend_consume_func_return_mrf32_mri64(%[[RET0]], %[[RET1]])
+// CHECK-SAME:          : (memref<*xf32>, memref<*xi64>) -> ()
+// CHECK:           return
+
+func @two_return_values(%arg0: memref<?xf32>, %arg1: memref<?xi64>) -> (memref<?xf32>, memref<?xi64>) {
+  return %arg0 ,%arg1 : memref<?xf32>, memref<?xi64>
+}