[mlir][Linalg] Employ finer-grained control of C interface emission

[nicolasvasilache]

Summary

Linalg makes it possible to interface codegen with externally precompiled HPC libraries. The mechanism to allow such interop uses a normalized ABI and the emission of C interface wrappers.

The mechanism controlling these C interface emission is too aggressive and makes it very easy to obtained undefined symbols for external function (e.g. the ones coming from libm).

This revision uses the newly introduced llvm.emit_c_interface function attribute which allows controlling this behavior at a function granularity. As a consequence LinalgToLLVM does not need to activate the C wrapper emission when adding the StdToLLVM patterns.

[mehdi_amini]

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

406

Can you add a comment explaining this?
This is obvious in the context of this diff, but I'm thinking about someone who does not know the detail of this dance and trying to read this code later.

[nicolasvasilache]

Add comment.

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

assert(op->getNumResults() == 0 &&
          "Library call for linalg operation can be generated only for ops that "
          "have void return types");
   auto libFnType = FunctionType::get(inputTypes, {}, rewriter.getContext());
 
   OpBuilder::InsertionGuard guard(rewriter);
   // Insert before module terminator.
   rewriter.setInsertionPoint(module.getBody(),
                              std::prev(module.getBody()->end()));
+  FuncOp funcOp =
       rewriter.create<FuncOp>(op->getLoc(), fnNameAttr.getValue(), libFnType,
                               ArrayRef<NamedAttribute>{});
+  // Insert a function attribute that will trigger the emission of the
+  // corresponding `_mlir_ciface_xxx` interface so that external libraries see
+  // a normalized ABI. This interface is added during std to llvm conversion.
+  funcOp.setAttr("llvm.emit_c_interface", UnitAttr::get(op->getContext()));
   return fnNameAttr;
 }
 
 namespace {
 
 // LinalgOpConversion<LinalgOp> creates a new call to the
 // `LinalgOp::getLibraryCallName()` function.
 // The implementation of the function can be either in the same module or in an

 void ConvertLinalgToLLVMPass::runOnModule() {
   auto module = getModule();
 
   // Convert to the LLVM IR dialect using the converter defined above.
   OwningRewritePatternList patterns;
   LLVMTypeConverter converter(&getContext());
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());
-  populateStdToLLVMConversionPatterns(converter, patterns, /*useAlloca=*/false,
-                                      /*emitCWrappers=*/true);
+  populateStdToLLVMConversionPatterns(converter, patterns);
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateLinalgToStandardConversionPatterns(patterns, &getContext());
   populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext());
 
   LLVMConversionTarget target(getContext());
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   target.addLegalOp<ModuleOp, ModuleTerminatorOp>();
   if (failed(applyFullConversion(module, target, patterns, &converter)))
     signalPassFailure();
 }
 
 std::unique_ptr<OpPassBase<ModuleOp>> mlir::createConvertLinalgToLLVMPass() {
   return std::make_unique<ConvertLinalgToLLVMPass>();
 }

mlir/test/mlir-cpu-runner/unranked_memref.mlir

func @main() -> () {
     store %i8_z, %I8[]: memref<i8>
     %U4 = memref_cast %I8 : memref<i8> to memref<*xi8>
     call @print_memref_i8(%U4) : (memref<*xi8>) -> ()
 
     dealloc %A : memref<10x3xf32, 0>
     return
 }
 
-func @print_memref_i8(memref<*xi8>)
-func @print_memref_f32(memref<*xf32>)
+func @print_memref_i8(memref<*xi8>) attributes { llvm.emit_c_interface }
+func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }

mlir/test/mlir-cpu-runner/utils.mlir

 }
 // PRINT-3D: Memref base@ = {{.*}} rank = 3 offset = 0 sizes = [3, 4, 5] strides = [20, 5, 1] data =
 // PRINT-3D-COUNT-4: {{.*[[:space:]].*}}2,    2,    2,    2,    2
 // PRINT-3D-COUNT-4: {{.*[[:space:]].*}}2,    2,    2,    2,    2
 // PRINT-3D-COUNT-2: {{.*[[:space:]].*}}2,    2,    2,    2,    2
 //    PRINT-3D-NEXT: 2,    2,    4,    2,    2
 //    PRINT-3D-NEXT: 2,    2,    2,    2,    2
 
-func @print_memref_f32(memref<*xf32>)
+func @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
 
 !vector_type_C = type vector<4x4xf32>
 !matrix_type_CC = type memref<1x1x!vector_type_C>
 func @vector_splat_2d() {
   %c0 = constant 0 : index
   %f10 = constant 10.0 : f32
   %vf10 = splat %f10: !vector_type_C
   %C = alloc() : !matrix_type_CC
   store %vf10, %C[%c0, %c0]: !matrix_type_CC
 
   %CC = memref_cast %C: !matrix_type_CC to memref<?x?x!vector_type_C>
   call @print_memref_vector_4x4xf32(%CC): (memref<?x?x!vector_type_C>) -> ()
 
   dealloc %C : !matrix_type_CC
   return
 }
 
 // PRINT-VECTOR-SPLAT-2D: Memref base@ = {{.*}} rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data =
 // PRINT-VECTOR-SPLAT-2D-NEXT: [((10, 10, 10, 10),   (10, 10, 10, 10),   (10, 10, 10, 10),   (10, 10, 10, 10))]
 
-func @print_memref_vector_4x4xf32(memref<?x?x!vector_type_C>)
+func @print_memref_vector_4x4xf32(memref<?x?x!vector_type_C>) attributes { llvm.emit_c_interface }