double check two returns from refbackend (#135)

Author: makslevental

mlir/extras/runtime/refbackend.py

@@ -179,22 +179,17 @@ def make_return_consumer(kernel_func):
 # This function will be called KERNEL_NAME_capi_wrapper and will have a {llvm.emit_c_interface} attribute.
 # Note there might be other such functions in the final module (gpu-lower-to-nvvm-pipeline somehow also inserts some like this).
 def make_kernel_wrapper(kernel_func, return_consumer=None):
-    c_api_compatible_types = [
-        T.memref(element_type=t.element_type) if MemRefType.isinstance(t) else t
-        for t in kernel_func.function_type.value.results
-    ]
-
     input_types = kernel_func.function_type.value.inputs
 
     @FuncOp.from_py_func(*input_types, name=f"{kernel_func.name.value}_capi_wrapper")
     def wrapper(*args, **_kwargs):
         results = CallOp(kernel_func, list(args)).results
-        c_api_compatible_results = []
-        for i, a in enumerate(results):
-            if MemRefType.isinstance(a.type):
-                a = cast(c_api_compatible_types[i], a)
-            c_api_compatible_results.append(a)
         if return_consumer is not None:
+            c_api_compatible_results = []
+            for i, a in enumerate(results):
+                if MemRefType.isinstance(a.type):
+                    a = cast(T.memref(element_type=a.type.element_type), a)
+                c_api_compatible_results.append(a)
             CallOp(return_consumer, c_api_compatible_results)
 
     wrapper_func_op = wrapper.func_op

tests/test_runtime.py

@@ -505,7 +505,7 @@ def memfoo(mem: ranked_memref_kxk_f32):
                 T.f32(), index_cast(T.i32(), i)
             )
             res = yield it_mem
-        return res
+        return res, res
 
     memfoo.emit()
 
@@ -520,8 +520,9 @@ def memfoo(mem: ranked_memref_kxk_f32):
     A = np.ones((K, K)).astype(np.float32)
     AA = ctypes.pointer(ctypes.pointer(get_ranked_memref_descriptor(A)))
 
-    results = invoker.memfoo_capi_wrapper(AA)
+    results, results1 = invoker.memfoo_capi_wrapper(AA)
     assert np.array_equal(np.diagonal(results), np.arange(1, K + 1))
+    assert np.array_equal(np.diagonal(results1), np.arange(1, K + 1))
 
 
 def test_setting_memref_diagonal_no_iter(ctx: MLIRContext, backend: LLVMJITBackend):
@@ -535,6 +536,7 @@ def memfoo(mem: ranked_memref_kxk_f32):
             mem[i, i] = mem[i, i] + mem[i, i] * sitofp(T.f32(), index_cast(T.i32(), i))
 
     memfoo.emit()
+    print(ctx.module)
 
     module = backend.compile(
         ctx.module,