[mlir][llvm] Pass struct results as parameter in c wrapper.

[herhut]

Summary

Returning structs directly in LLVM does not necessarily align with the C ABI of the platform. This might happen to work on Linux but for small structs this breaks on Windows. With this change, the wrappers work platform independently.

[ftynse]

mlir/docs/LLVMDialectMemRefConvention.md

235

s/reference/pointer

399–404

It's already defined above

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1319

Nit: just cast

[herhut]

Review comments...

[herhut]

Thanks!

mlir/docs/LLVMDialectMemRefConvention.md

399–404

Copy/paste mistake. Thanks!

mlir/docs/LLVMDialectMemRefConvention.md

struct MemRefDescriptor {
   T *allocated;
   T *aligned;
   intptr_t offset;
   intptr_t sizes[N];
   intptr_t strides[N];
 };
 ```
 
+Furthermore, we also rewrite function results to pointer parameters if the
+rewritten function result has a struct type. The special result parameter is
+added as the first parameter and is of pointer-to-struct type.
+
 If enabled, the option will do the following. For _external_ functions declared
 in the MLIR module.
 
 1.  Declare a new function `_mlir_ciface_<original name>` where memref arguments
     are converted to pointer-to-struct and the remaining arguments are converted
-    as usual.
-1.  Add a body to the original function (making it non-external) that
-    1.  allocates a memref descriptor,
-    1.  populates it, and
-    1.  passes the pointer to it into the newly declared interface function,
+    as usual. Results are converted to a special argument if they are of struct
+    type.
+2.  Add a body to the original function (making it non-external) that
+    1.  allocates memref descriptors,
+    2.  populates them,
+    3.  potentially allocates space for the result struct, and
+    4.  passes the pointers to these into the newly declared interface function,
         then
-    1.  collects the result of the call and returns it to the caller.
+    5.  collects the result of the call (potentially from the result struct),
+        and
+    6.  returns it to the caller.
 
 For (non-external) functions defined in the MLIR module.
 
 1.  Define a new function `_mlir_ciface_<original name>` where memref arguments
     are converted to pointer-to-struct and the remaining arguments are converted
-    as usual.
-1.  Populate the body of the newly defined function with IR that
+    as usual. Results are converted to a special argument if they are of struct
+    type.
+2.  Populate the body of the newly defined function with IR that
     1.  loads descriptors from pointers;
-    1.  unpacks descriptor into individual non-aggregate values;
-    1.  passes these values into the original function;
-    1.  collects the result of the call and returns it to the caller.
+    2.  unpacks descriptor into individual non-aggregate values;
+    3.  passes these values into the original function;
+    4.  collects the results of the call and
+    5.  either copies the results into the result struct or returns them to the
+        caller.
 
 Examples:
 
 ```mlir
 
 func @qux(%arg0: memref<?x?xf32>)
 
 // Gets converted into the following
llvm.func @_mlir_ciface_foo(%arg0: !llvm.memref_2d_ptr) {
   %7 = llvm.extractvalue %0[4, 1] : !llvm.memref_2d
   llvm.call @foo(%1, %2, %3, %4, %5, %6, %7)
     : (!llvm.ptr<f32>, !llvm.ptr<f32>, i64, i64, i64,
        i64, i64) -> ()
   llvm.return
 }
 ```
 
+```mlir
+func @foo(%arg0: memref<?x?xf32>) -> memref<?x?xf32> {
+  return %arg0 : memref<?x?xf32>
+}
+
+// Gets converted into the following
+// (using type alias for brevity):
+!llvm.memref_2d = type !llvm.struct<(ptr<f32>, ptr<f32>, i64,
+                                     array<2xi64>, array<2xi64>)>
+!llvm.memref_2d_ptr = type !llvm.ptr<struct<(ptr<f32>, ptr<f32>, i64,
+                                             array<2xi64>, array<2xi64>)>>
+
+// Function with unpacked arguments.
+llvm.func @foo(%arg0: !llvm.ptr<f32>, %arg1: !llvm.ptr<f32>, %arg2: i64,
+               %arg3: i64, %arg4: i64, %arg5: i64, %arg6: i64)
+    -> !llvm.memref_2d {
+  %0 = llvm.mlir.undef : !llvm.memref_2d
+  %1 = llvm.insertvalue %arg0, %0[0] : !llvm.memref_2d
+  %2 = llvm.insertvalue %arg1, %1[1] : !llvm.memref_2d
+  %3 = llvm.insertvalue %arg2, %2[2] : !llvm.memref_2d
+  %4 = llvm.insertvalue %arg3, %3[3, 0] : !llvm.memref_2d
+  %5 = llvm.insertvalue %arg5, %4[4, 0] : !llvm.memref_2d
+  %6 = llvm.insertvalue %arg4, %5[3, 1] : !llvm.memref_2d
+  %7 = llvm.insertvalue %arg6, %6[4, 1] : !llvm.memref_2d
+  llvm.return %7 : !llvm.memref_2d
+}
+
+// Interface function callable from C.
+llvm.func @_mlir_ciface_foo(%arg0: !llvm.memref_2d_ptr, %arg1: !llvm.memref_2d_ptr) {
+  %0 = llvm.load %arg1 : !llvm.memref_2d_ptr
+  %1 = llvm.extractvalue %0[0] : !llvm.memref_2d
+  %2 = llvm.extractvalue %0[1] : !llvm.memref_2d
+  %3 = llvm.extractvalue %0[2] : !llvm.memref_2d
+  %4 = llvm.extractvalue %0[3, 0] : !llvm.memref_2d
+  %5 = llvm.extractvalue %0[3, 1] : !llvm.memref_2d
+  %6 = llvm.extractvalue %0[4, 0] : !llvm.memref_2d
+  %7 = llvm.extractvalue %0[4, 1] : !llvm.memref_2d
+  %8 = llvm.call @foo(%1, %2, %3, %4, %5, %6, %7)
+    : (!llvm.ptr<f32>, !llvm.ptr<f32>, i64, i64, i64, i64, i64) -> !llvm.memref_2d
+  llvm.store %8, %arg0 : !llvm.memref_2d_ptr
+  llvm.return
+}
+
 Rationale: Introducing auxiliary functions for C-compatible interfaces is
 preferred to modifying the calling convention since it will minimize the effect
 of C compatibility on intra-module calls or calls between MLIR-generated
 functions. In particular, when calling external functions from an MLIR module in
 a (parallel) loop, the fact of storing a memref descriptor on stack can lead to
 stack exhaustion and/or concurrent access to the same address. Auxiliary
 interface function serves as an allocation scope in this case. Furthermore, when
 targeting accelerators with separate memory spaces such as GPUs, stack-allocated
 descriptors passed by pointer would have to be transferred to the device memory,
 which introduces significant overhead. In such situations, auxiliary interface
 functions are executed on host and only pass the values through device function
 invocation mechanism.
 
 ## Default Memref Model

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

public:
 
   /// Promote the LLVM struct representation of one MemRef descriptor to stack
   /// and use pointer to struct to avoid the complexity of the platform-specific
   /// C/C++ ABI lowering related to struct argument passing.
   Value promoteOneMemRefDescriptor(Location loc, Value operand,
                                    OpBuilder &builder);
 
   /// Converts the function type to a C-compatible format, in particular using
-  /// pointers to memref descriptors for arguments.
-  Type convertFunctionTypeCWrapper(FunctionType type);
+  /// pointers to memref descriptors for arguments. Also converts the return
+  /// type to a pointer argument if it is a struct. Returns true if this
+  /// was the case.
+  std::pair<Type, bool> convertFunctionTypeCWrapper(FunctionType type);
 
   /// Returns the data layout to use during and after conversion.
   const llvm::DataLayout &getDataLayout() { return options.dataLayout; }
 
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   Type getIndexType();
 

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

Type resultType = funcTy.getNumResults() == 0
                         : unwrap(packFunctionResults(funcTy.getResults()));
   if (!resultType)
     return {};
   return LLVM::LLVMFunctionType::get(resultType, argTypes, isVariadic);
 }
 
 /// Converts the function type to a C-compatible format, in particular using
 /// pointers to memref descriptors for arguments.
-Type LLVMTypeConverter::convertFunctionTypeCWrapper(FunctionType type) {
+std::pair<Type, bool>
+LLVMTypeConverter::convertFunctionTypeCWrapper(FunctionType type) {
   SmallVector<Type, 4> inputs;
+  bool resultIsNowArg = false;
+
+  Type resultType = type.getNumResults() == 0
+                        ? LLVM::LLVMVoidType::get(&getContext())
+                        : unwrap(packFunctionResults(type.getResults()));
+  if (!resultType)
+    return {};
+
+  if (auto structType = resultType.dyn_cast<LLVM::LLVMStructType>()) {
+    // Struct types cannot be safely returned via C interface. Make this a
+    // pointer argument, instead.
+    inputs.push_back(LLVM::LLVMPointerType::get(structType));
+    resultType = LLVM::LLVMVoidType::get(&getContext());
+    resultIsNowArg = true;
+  }
 
   for (Type t : type.getInputs()) {
     auto converted = convertType(t);
     if (!converted || !LLVM::isCompatibleType(converted))
       return {};
     if (t.isa<MemRefType, UnrankedMemRefType>())
       converted = LLVM::LLVMPointerType::get(converted);
     inputs.push_back(converted);
   }
 
-  Type resultType = type.getNumResults() == 0
-                        ? LLVM::LLVMVoidType::get(&getContext())
-                        : unwrap(packFunctionResults(type.getResults()));
-  if (!resultType)
-    return {};
-
-  return LLVM::LLVMFunctionType::get(resultType, inputs);
+  return {LLVM::LLVMFunctionType::get(resultType, inputs), resultIsNowArg};
 }
 
 static constexpr unsigned kAllocatedPtrPosInMemRefDescriptor = 0;
 static constexpr unsigned kAlignedPtrPosInMemRefDescriptor = 1;
 static constexpr unsigned kOffsetPosInMemRefDescriptor = 2;
 static constexpr unsigned kSizePosInMemRefDescriptor = 3;
 static constexpr unsigned kStridePosInMemRefDescriptor = 4;
 
if (attr.first == SymbolTable::getSymbolAttrName() ||
       continue;
     result.push_back(attr);
   }
 }
 
 /// Creates an auxiliary function with pointer-to-memref-descriptor-struct
 /// arguments instead of unpacked arguments. This function can be called from C
 /// by passing a pointer to a C struct corresponding to a memref descriptor.
+/// Similarly, returned memrefs are passed via pointers to a C struct that is
+/// passed as additional argument.
 /// Internally, the auxiliary function unpacks the descriptor into individual
-/// components and forwards them to `newFuncOp`.
+/// components and forwards them to `newFuncOp` and forwards the results to
+/// the extra arguments.
 static void wrapForExternalCallers(OpBuilder &rewriter, Location loc,
                                    LLVMTypeConverter &typeConverter,
                                    FuncOp funcOp, LLVM::LLVMFuncOp newFuncOp) {
   auto type = funcOp.getType();
   SmallVector<NamedAttribute, 4> attributes;
   filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/false,
                        attributes);
+  Type wrapperFuncType;
+  bool resultIsNowArg;
+  std::tie(wrapperFuncType, resultIsNowArg) =
+      typeConverter.convertFunctionTypeCWrapper(type);
   auto wrapperFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
       loc, llvm::formatv("_mlir_ciface_{0}", funcOp.getName()).str(),
-      typeConverter.convertFunctionTypeCWrapper(type), LLVM::Linkage::External,
-      attributes);
+      wrapperFuncType, LLVM::Linkage::External, attributes);
 
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPointToStart(wrapperFuncOp.addEntryBlock());
 
   SmallVector<Value, 8> args;
+  size_t argOffset = resultIsNowArg ? 1 : 0;
   for (auto &en : llvm::enumerate(type.getInputs())) {
-    Value arg = wrapperFuncOp.getArgument(en.index());
+    Value arg = wrapperFuncOp.getArgument(en.index() + argOffset);
     if (auto memrefType = en.value().dyn_cast<MemRefType>()) {
       Value loaded = rewriter.create<LLVM::LoadOp>(loc, arg);
       MemRefDescriptor::unpack(rewriter, loc, loaded, memrefType, args);
       continue;
     }
     if (en.value().isa<UnrankedMemRefType>()) {
       Value loaded = rewriter.create<LLVM::LoadOp>(loc, arg);
       UnrankedMemRefDescriptor::unpack(rewriter, loc, loaded, args);
       continue;
     }
 
-    args.push_back(wrapperFuncOp.getArgument(en.index()));
+    args.push_back(arg);
   }
+
   auto call = rewriter.create<LLVM::CallOp>(loc, newFuncOp, args);
+
+  if (resultIsNowArg) {
+    rewriter.create<LLVM::StoreOp>(loc, call.getResult(0),
+                                   wrapperFuncOp.getArgument(0));
+    rewriter.create<LLVM::ReturnOp>(loc, ValueRange{});
+  } else {
     rewriter.create<LLVM::ReturnOp>(loc, call.getResults());
   }
+}
 
 /// Creates an auxiliary function with pointer-to-memref-descriptor-struct
 /// arguments instead of unpacked arguments. Creates a body for the (external)
 /// `newFuncOp` that allocates a memref descriptor on stack, packs the
 /// individual arguments into this descriptor and passes a pointer to it into
-/// the auxiliary function. This auxiliary external function is now compatible
-/// with functions defined in C using pointers to C structs corresponding to a
-/// memref descriptor.
+/// the auxiliary function. If the result of the function cannot be directly
+/// returned, we write it to a special first argument that provides a pointer
+/// to a corresponding struct. This auxiliary external function is now
+/// compatible with functions defined in C using pointers to C structs
+/// corresponding to a memref descriptor.
 static void wrapExternalFunction(OpBuilder &builder, Location loc,
                                  LLVMTypeConverter &typeConverter,
                                  FuncOp funcOp, LLVM::LLVMFuncOp newFuncOp) {
   OpBuilder::InsertionGuard guard(builder);
 
-  Type wrapperType =
+  Type wrapperType;
+  bool resultIsNowArg;
+  std::tie(wrapperType, resultIsNowArg) =
       typeConverter.convertFunctionTypeCWrapper(funcOp.getType());
   // This conversion can only fail if it could not convert one of the argument
-  // types. But since it has been applies to a non-wrapper function before, it
+  // types. But since it has been applied to a non-wrapper function before, it
   // should have failed earlier and not reach this point at all.
   assert(wrapperType && "unexpected type conversion failure");
 
   SmallVector<NamedAttribute, 4> attributes;
   filterFuncAttributes(funcOp->getAttrs(), /*filterArgAttrs=*/false,
                        attributes);
 
   // Create the auxiliary function.
   auto wrapperFunc = builder.create<LLVM::LLVMFuncOp>(
       loc, llvm::formatv("_mlir_ciface_{0}", funcOp.getName()).str(),
       wrapperType, LLVM::Linkage::External, attributes);
 
   builder.setInsertionPointToStart(newFuncOp.addEntryBlock());
 
   // Get a ValueRange containing arguments.
   FunctionType type = funcOp.getType();
   SmallVector<Value, 8> args;
   args.reserve(type.getNumInputs());
   ValueRange wrapperArgsRange(newFuncOp.getArguments());
 
+  if (resultIsNowArg) {
+    // Allocate the struct on the stack and pass the pointer.
+    Type resultType =
+        wrapperType.cast<LLVM::LLVMFunctionType>().getParamType(0);
+    Value one = builder.create<LLVM::ConstantOp>(
+        loc, typeConverter.convertType(builder.getIndexType()),
+        builder.getIntegerAttr(builder.getIndexType(), 1));
+    Value result = builder.create<LLVM::AllocaOp>(loc, resultType, one);
+    args.push_back(result);
+  }
+
   // Iterate over the inputs of the original function and pack values into
   // memref descriptors if the original type is a memref.
   for (auto &en : llvm::enumerate(type.getInputs())) {
     Value arg;
     int numToDrop = 1;
     auto memRefType = en.value().dyn_cast<MemRefType>();
     auto unrankedMemRefType = en.value().dyn_cast<UnrankedMemRefType>();
     if (memRefType || unrankedMemRefType) {
for (auto &en : llvm::enumerate(type.getInputs())) {
     }
 
     args.push_back(arg);
     wrapperArgsRange = wrapperArgsRange.drop_front(numToDrop);
   }
   assert(wrapperArgsRange.empty() && "did not map some of the arguments");
 
   auto call = builder.create<LLVM::CallOp>(loc, wrapperFunc, args);
+
+  if (resultIsNowArg) {
+    Value result = builder.create<LLVM::LoadOp>(loc, args.front());
+    builder.create<LLVM::ReturnOp>(loc, ValueRange{result});
+  } else {
     builder.create<LLVM::ReturnOp>(loc, call.getResults());
   }
+}
 
 namespace {
 
 struct FuncOpConversionBase : public ConvertOpToLLVMPattern<FuncOp> {
 protected:
   using ConvertOpToLLVMPattern<FuncOp>::ConvertOpToLLVMPattern;
 
   // Convert input FuncOp to LLVMFuncOp by using the LLVMTypeConverter provided

mlir/test/Conversion/StandardToLLVM/calling-convention.mlir

func @return_var_memref_caller(%arg0: memref<4x3xf32>) {
   // CHECK: %[[DESC:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[CALL_RES]][0] : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[DESC_1:.*]] = llvm.insertvalue %[[RANK]], %[[DESC]][0]
   // CHECK: llvm.insertvalue %[[ALLOCA]], %[[DESC_1]][1]
   return
 }
 
 // CHECK-LABEL: llvm.func @return_var_memref
-func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> {
+func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> attributes { llvm.emit_c_interface } {
   // Match the construction of the unranked descriptor.
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca
   // CHECK: %[[MEMORY:.*]] = llvm.bitcast %[[ALLOCA]]
   // CHECK: %[[DESC_0:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[DESC_1:.*]] = llvm.insertvalue %{{.*}}, %[[DESC_0]][0]
   // CHECK: %[[DESC_2:.*]] = llvm.insertvalue %[[MEMORY]], %[[DESC_1]][1]
   %0 = memref_cast %arg0: memref<4x3xf32> to memref<*xf32>
 
func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> attributes { llvm.emit_c_interface } {
   // CHECK: %[[NEW_DESC:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[DESC_2]][0] : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[NEW_DESC_1:.*]] = llvm.insertvalue %[[RANK]], %[[NEW_DESC]][0]
   // CHECK: %[[NEW_DESC_2:.*]] = llvm.insertvalue %[[ALLOCATED]], %[[NEW_DESC_1]][1]
   // CHECK: llvm.return %[[NEW_DESC_2]]
   return %0 : memref<*xf32>
 }
 
+// Check that the result memref is passed as parameter
+// CHECK-LABEL: @_mlir_ciface_return_var_memref
+// CHECK-SAME: (%{{.*}}: !llvm.ptr<struct<(i64, ptr<i8>)>>, %{{.*}}: !llvm.ptr<struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>>)
+
 // CHECK-LABEL: llvm.func @return_two_var_memref_caller
 func @return_two_var_memref_caller(%arg0: memref<4x3xf32>) {
   // Only check that we create two different descriptors using different
   // memory, and deallocate both sources. The size computation is same as for
   // the single result.
   // CHECK: %[[CALL_RES:.*]] = llvm.call @return_two_var_memref
   // CHECK: %[[RES_1:.*]] = llvm.extractvalue %[[CALL_RES]][0]
   // CHECK: %[[RES_2:.*]] = llvm.extractvalue %[[CALL_RES]][1]
func @return_two_var_memref_caller(%arg0: memref<4x3xf32>) {
   // CHECK: llvm.call @free(%[[SOURCE_2]])
   // CHECK: %[[DESC_2:.*]] = llvm.mlir.undef : ![[DESC_TYPE]]
   // CHECK: %[[DESC_21:.*]] = llvm.insertvalue %{{.*}}, %[[DESC_2]][0]
   // CHECK: llvm.insertvalue %[[ALLOCA_2]], %[[DESC_21]][1]
   return
 }
 
 // CHECK-LABEL: llvm.func @return_two_var_memref
-func @return_two_var_memref(%arg0: memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) {
+func @return_two_var_memref(%arg0: memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) attributes { llvm.emit_c_interface } {
   // Match the construction of the unranked descriptor.
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca
   // CHECK: %[[MEMORY:.*]] = llvm.bitcast %[[ALLOCA]]
   // CHECK: %[[DESC_0:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[DESC_1:.*]] = llvm.insertvalue %{{.*}}, %[[DESC_0]][0]
   // CHECK: %[[DESC_2:.*]] = llvm.insertvalue %[[MEMORY]], %[[DESC_1]][1]
   %0 = memref_cast %arg0 : memref<4x3xf32> to memref<*xf32>
 
func @return_two_var_memref(%arg0: memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) attributes { llvm.emit_c_interface } {
 
   // CHECK: %[[RESULTS:.*]] = llvm.mlir.undef : !llvm.struct<(struct<(i64, ptr<i8>)>, struct<(i64, ptr<i8>)>)>
   // CHECK: %[[RESULTS_1:.*]] = llvm.insertvalue %[[RES_12]], %[[RESULTS]]
   // CHECK: %[[RESULTS_2:.*]] = llvm.insertvalue %[[RES_22]], %[[RESULTS_1]]
   // CHECK: llvm.return %[[RESULTS_2]]
   return %0, %0 : memref<*xf32>, memref<*xf32>
 }
 
+// Check that the result memrefs are passed as parameter
+// CHECK-LABEL: @_mlir_ciface_return_two_var_memref
+// CHECK-SAME: (%{{.*}}: !llvm.ptr<struct<(struct<(i64, ptr<i8>)>, struct<(i64, ptr<i8>)>)>>,
+// CHECK-SAME: %{{.*}}: !llvm.ptr<struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>>)
+