[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 reference 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
+}
+
+// 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.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

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.dyn_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>)>>)
+