[mlir][LLVM] NFC - Refactor a lookupOrCreateFn to reuse common function creation.

[ftynse]

I believe async lowering also needs this, but leaving for future work.

mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h

10

among

56

Also add funcs for malloc and free?

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

39

Nit: spurious empty line here and below

85–88

Nit: could we make something like b.getSymbolRefAttr(fn) into a variable to make this formatting look less horrible?

[nicolasvasilache]

Address review.

[nicolasvasilache]

Thanks!

[nicolasvasilache]

Simplify constant exprs.

mlir/include/mlir/Dialect/LLVMIR/FunctionCallUtils.h

+//===- FunctionCallUtils.h - Utilities for C function calls -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares helper functions to call common simple C functions in
+// LLVMIR (e.g. amon others to support printing and debugging).
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LLVMIR_FUNCTIONCALLUTILS_H_
+#define MLIR_DIALECT_LLVMIR_FUNCTIONCALLUTILS_H_
+
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class Location;
+class ModuleOp;
+class OpBuilder;
+class Operation;
+class Type;
+class ValueRange;
+
+namespace LLVM {
+class LLVMFuncOp;
+
+struct CFunctionNames {
+  static constexpr llvm::StringRef kPrintI64 = "printI64";
+  static constexpr llvm::StringRef kPrintU64 = "printU64";
+  static constexpr llvm::StringRef kPrintF32 = "printF32";
+  static constexpr llvm::StringRef kPrintF64 = "printF64";
+  static constexpr llvm::StringRef kPrintOpen = "printOpen";
+  static constexpr llvm::StringRef kPrintClose = "printClose";
+  static constexpr llvm::StringRef kPrintComma = "printComma";
+  static constexpr llvm::StringRef kPrintNewline = "printNewline";
+  static constexpr llvm::StringRef kMalloc = "malloc";
+  static constexpr llvm::StringRef kAlignedAlloc = "aligned_alloc";
+};
+
+/// Helper functions to lookup or create the declaration for commonly used
+/// external C function calls. Such ops can then be invoked by creating a CallOp
+/// with the proper arguments via `createLLVMCall`.
+/// The list of functions provided here must be implemented separately (e.g. as
+/// part of a support runtime library or as part of the libc).
+LLVM::LLVMFuncOp lookupOrCreatePrintI64Fn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintU64Fn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintF32Fn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintF64Fn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintOpenFn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintCloseFn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintCommaFn(ModuleOp moduleOp);
+LLVM::LLVMFuncOp lookupOrCreatePrintNewlineFn(ModuleOp moduleOp);
+
+/// Create a FuncOp with signature `resultTyp`(`paramTypes`)` and name `name`.
+LLVM::LLVMFuncOp lookupOrCreateFn(ModuleOp moduleOp, StringRef name,
+                                  ArrayRef<Type> paramTypes = {},
+                                  Type resultType = {});
+
+/// Helper wrapper to create a call to `fn` with `args` and `resultTypes`.
+Operation::result_range createLLVMCall(OpBuilder &b, Location loc,
+                                       LLVM::LLVMFuncOp fn,
+                                       ValueRange args = {},
+                                       ArrayRef<Type> resultTypes = {});
+
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LLVMIR_FUNCTIONCALLUTILS_H_

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //===- StandardToLLVM.cpp - Standard to LLVM dialect conversion -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to convert MLIR standard and builtin dialects
 // into the LLVM IR dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
static Value createAligned(ConversionPatternRewriter &rewriter, Location loc,
                              Value input, Value alignment) {
     Value one = createIndexAttrConstant(rewriter, loc, alignment.getType(), 1);
     Value bump = rewriter.create<LLVM::SubOp>(loc, alignment, one);
     Value bumped = rewriter.create<LLVM::AddOp>(loc, input, bump);
     Value mod = rewriter.create<LLVM::URemOp>(loc, bumped, alignment);
     return rewriter.create<LLVM::SubOp>(loc, bumped, mod);
   }
 
-  // Creates a call to an allocation function with params and casts the
-  // resulting void pointer to ptrType.
-  Value createAllocCall(Location loc, StringRef name, Type ptrType,
-                        ArrayRef<Value> params, ModuleOp module,
-                        ConversionPatternRewriter &rewriter) const {
-    SmallVector<Type, 2> paramTypes;
-    auto allocFuncOp = module.lookupSymbol<LLVM::LLVMFuncOp>(name);
-    if (!allocFuncOp) {
-      for (Value param : params)
-        paramTypes.push_back(param.getType());
-      auto allocFuncType =
-          LLVM::LLVMFunctionType::get(getVoidPtrType(), paramTypes);
-      OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.setInsertionPointToStart(module.getBody());
-      allocFuncOp = rewriter.create<LLVM::LLVMFuncOp>(rewriter.getUnknownLoc(),
-                                                      name, allocFuncType);
-    }
-    auto allocFuncSymbol = rewriter.getSymbolRefAttr(allocFuncOp);
-    auto allocatedPtr = rewriter
-                            .create<LLVM::CallOp>(loc, getVoidPtrType(),
-                                                  allocFuncSymbol, params)
-                            .getResult(0);
-    return rewriter.create<LLVM::BitcastOp>(loc, ptrType, allocatedPtr);
-  }
-
   /// Allocates the underlying buffer. Returns the allocated pointer and the
   /// aligned pointer.
   virtual std::tuple<Value, Value>
   allocateBuffer(ConversionPatternRewriter &rewriter, Location loc,
                  Value sizeBytes, Operation *op) const = 0;
 
 private:
   static MemRefType getMemRefResultType(Operation *op) {
std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
     if (alignment) {
       // Adjust the allocation size to consider alignment.
       sizeBytes = rewriter.create<LLVM::AddOp>(loc, sizeBytes, alignment);
     }
 
     // Allocate the underlying buffer and store a pointer to it in the MemRef
     // descriptor.
     Type elementPtrType = this->getElementPtrType(memRefType);
+    auto allocFuncOp = LLVM::lookupOrCreateFn(
+        allocOp->getParentOfType<ModuleOp>(), LLVM::CFunctionNames::kMalloc,
+        {getIndexType()}, getVoidPtrType());
+    auto results = createLLVMCall(rewriter, loc, allocFuncOp, {sizeBytes},
+                                  getVoidPtrType());
     Value allocatedPtr =
-        createAllocCall(loc, "malloc", elementPtrType, {sizeBytes},
-                        allocOp->getParentOfType<ModuleOp>(), rewriter);
+        rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     Value alignedPtr = allocatedPtr;
     if (alignment) {
       auto intPtrType = getIntPtrType(memRefType.getMemorySpace());
       // Compute the aligned type pointer.
       Value allocatedInt =
           rewriter.create<LLVM::PtrToIntOp>(loc, intPtrType, allocatedPtr);
       Value alignmentInt =
std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
     Value allocAlignment = createIndexConstant(rewriter, loc, alignment);
 
     // aligned_alloc requires size to be a multiple of alignment; we will pad
     // the size to the next multiple if necessary.
     if (!isMemRefSizeMultipleOf(memRefType, alignment))
       sizeBytes = createAligned(rewriter, loc, sizeBytes, allocAlignment);
 
     Type elementPtrType = this->getElementPtrType(memRefType);
-    Value allocatedPtr = createAllocCall(
-        loc, "aligned_alloc", elementPtrType, {allocAlignment, sizeBytes},
-        allocOp->getParentOfType<ModuleOp>(), rewriter);
+    auto allocFuncOp = LLVM::lookupOrCreateFn(
+        allocOp->getParentOfType<ModuleOp>(),
+        LLVM::CFunctionNames::kAlignedAlloc, {getIndexType(), getIndexType()},
+        getVoidPtrType());
+    auto results =
+        createLLVMCall(rewriter, loc, allocFuncOp, {allocAlignment, sizeBytes},
+                       getVoidPtrType());
+    Value allocatedPtr =
+        rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     return std::make_tuple(allocatedPtr, allocatedPtr);
   }
 
   /// The minimum alignment to use with aligned_alloc (has to be a power of 2).
   static constexpr uint64_t kMinAlignedAllocAlignment = 16UL;
 };
 

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Target/LLVMIR/TypeTranslation.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
if (typeConverter->convertType(printType) == nullptr)
       return failure();
 
     // Make sure element type has runtime support.
     PrintConversion conversion = PrintConversion::None;
     VectorType vectorType = printType.dyn_cast<VectorType>();
     Type eltType = vectorType ? vectorType.getElementType() : printType;
     Operation *printer;
     if (eltType.isF32()) {
-      printer = getPrintFloat(printOp);
+      printer =
+          LLVM::lookupOrCreatePrintF32Fn(printOp->getParentOfType<ModuleOp>());
     } else if (eltType.isF64()) {
-      printer = getPrintDouble(printOp);
+      printer =
+          LLVM::lookupOrCreatePrintF64Fn(printOp->getParentOfType<ModuleOp>());
     } else if (eltType.isIndex()) {
-      printer = getPrintU64(printOp);
+      printer =
+          LLVM::lookupOrCreatePrintU64Fn(printOp->getParentOfType<ModuleOp>());
     } else if (auto intTy = eltType.dyn_cast<IntegerType>()) {
       // Integers need a zero or sign extension on the operand
       // (depending on the source type) as well as a signed or
       // unsigned print method. Up to 64-bit is supported.
       unsigned width = intTy.getWidth();
       if (intTy.isUnsigned()) {
         if (width <= 64) {
           if (width < 64)
             conversion = PrintConversion::ZeroExt64;
-          printer = getPrintU64(printOp);
+          printer = LLVM::lookupOrCreatePrintU64Fn(
+              printOp->getParentOfType<ModuleOp>());
         } else {
           return failure();
         }
       } else {
         assert(intTy.isSignless() || intTy.isSigned());
         if (width <= 64) {
           // Note that we *always* zero extend booleans (1-bit integers),
           // so that true/false is printed as 1/0 rather than -1/0.
           if (width == 1)
             conversion = PrintConversion::ZeroExt64;
           else if (width < 64)
             conversion = PrintConversion::SignExt64;
-          printer = getPrintI64(printOp);
+          printer = LLVM::lookupOrCreatePrintI64Fn(
+              printOp->getParentOfType<ModuleOp>());
         } else {
           return failure();
         }
       }
     } else {
       return failure();
     }
 
     // Unroll vector into elementary print calls.
     int64_t rank = vectorType ? vectorType.getRank() : 0;
     emitRanks(rewriter, printOp, adaptor.source(), vectorType, printer, rank,
               conversion);
-    emitCall(rewriter, printOp->getLoc(), getPrintNewline(printOp));
+    emitCall(rewriter, printOp->getLoc(),
+             LLVM::lookupOrCreatePrintNewlineFn(
+                 printOp->getParentOfType<ModuleOp>()));
     rewriter.eraseOp(printOp);
     return success();
   }
 
 private:
   enum class PrintConversion {
     // clang-format off
     None,
if (rank == 0) {
         break;
       case PrintConversion::None:
         break;
       }
       emitCall(rewriter, loc, printer, value);
       return;
     }
 
-    emitCall(rewriter, loc, getPrintOpen(op));
-    Operation *printComma = getPrintComma(op);
+    emitCall(rewriter, loc,
+             LLVM::lookupOrCreatePrintOpenFn(op->getParentOfType<ModuleOp>()));
+    Operation *printComma =
+        LLVM::lookupOrCreatePrintCommaFn(op->getParentOfType<ModuleOp>());
     int64_t dim = vectorType.getDimSize(0);
     for (int64_t d = 0; d < dim; ++d) {
       auto reducedType =
           rank > 1 ? reducedVectorTypeFront(vectorType) : nullptr;
       auto llvmType = typeConverter->convertType(
           rank > 1 ? reducedType : vectorType.getElementType());
       Value nestedVal = extractOne(rewriter, *getTypeConverter(), loc, value,
                                    llvmType, rank, d);
       emitRanks(rewriter, op, nestedVal, reducedType, printer, rank - 1,
                 conversion);
       if (d != dim - 1)
         emitCall(rewriter, loc, printComma);
     }
-    emitCall(rewriter, loc, getPrintClose(op));
+    emitCall(rewriter, loc,
+             LLVM::lookupOrCreatePrintCloseFn(op->getParentOfType<ModuleOp>()));
   }
 
   // Helper to emit a call.
   static void emitCall(ConversionPatternRewriter &rewriter, Location loc,
                        Operation *ref, ValueRange params = ValueRange()) {
     rewriter.create<LLVM::CallOp>(loc, TypeRange(),
                                   rewriter.getSymbolRefAttr(ref), params);
   }
-
-  // Helper for printer method declaration (first hit) and lookup.
-  static Operation *getPrint(Operation *op, StringRef name,
-                             ArrayRef<Type> params) {
-    auto module = op->getParentOfType<ModuleOp>();
-    auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(name);
-    if (func)
-      return func;
-    OpBuilder moduleBuilder(module.getBodyRegion());
-    return moduleBuilder.create<LLVM::LLVMFuncOp>(
-        op->getLoc(), name,
-        LLVM::LLVMFunctionType::get(LLVM::LLVMVoidType::get(op->getContext()),
-                                    params));
-  }
-
-  // Helpers for method names.
-  Operation *getPrintI64(Operation *op) const {
-    return getPrint(op, "printI64", IntegerType::get(op->getContext(), 64));
-  }
-  Operation *getPrintU64(Operation *op) const {
-    return getPrint(op, "printU64", IntegerType::get(op->getContext(), 64));
-  }
-  Operation *getPrintFloat(Operation *op) const {
-    return getPrint(op, "printF32", Float32Type::get(op->getContext()));
-  }
-  Operation *getPrintDouble(Operation *op) const {
-    return getPrint(op, "printF64", Float64Type::get(op->getContext()));
-  }
-  Operation *getPrintOpen(Operation *op) const {
-    return getPrint(op, "printOpen", {});
-  }
-  Operation *getPrintClose(Operation *op) const {
-    return getPrint(op, "printClose", {});
-  }
-  Operation *getPrintComma(Operation *op) const {
-    return getPrint(op, "printComma", {});
-  }
-  Operation *getPrintNewline(Operation *op) const {
-    return getPrint(op, "printNewline", {});
-  }
 };
 
 /// Progressive lowering of ExtractStridedSliceOp to either:
 ///   1. express single offset extract as a direct shuffle.
 ///   2. extract + lower rank strided_slice + insert for the n-D case.
 class VectorExtractStridedSliceOpConversion
     : public OpRewritePattern<ExtractStridedSliceOp> {
 public:

mlir/lib/Dialect/LLVMIR/IR/FunctionCallUtils.cpp

+//===- FunctionCallUtils.cpp - Utilities for C function calls -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements helper functions to call common simple C functions in
+// LLVMIR (e.g. amon others to support printing and debugging).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Support/LLVM.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+
+/// Generic print function lookupOrCreate helper.
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreateFn(ModuleOp moduleOp, StringRef name,
+                                              ArrayRef<Type> paramTypes,
+                                              Type resultType) {
+  auto func = moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(name);
+  if (func)
+    return func;
+  OpBuilder b(moduleOp.getBodyRegion());
+  return b.create<LLVM::LLVMFuncOp>(
+      moduleOp->getLoc(), name,
+      LLVM::LLVMFunctionType::get(resultType, paramTypes));
+}
+
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintI64Fn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintI64,
+                          IntegerType::get(moduleOp->getContext(), 64),
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintU64Fn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintU64,
+                          IntegerType::get(moduleOp->getContext(), 64),
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF32Fn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintF32,
+                          Float32Type::get(moduleOp->getContext()),
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintF64Fn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintF64,
+                          Float64Type::get(moduleOp->getContext()),
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintOpenFn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintOpen, {},
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintCloseFn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintClose, {},
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintCommaFn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintComma, {},
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintNewlineFn(ModuleOp moduleOp) {
+  return lookupOrCreateFn(moduleOp, CFunctionNames::kPrintNewline, {},
+
+                          LLVM::LLVMVoidType::get(moduleOp->getContext()));
+}
+
+Operation::result_range mlir::LLVM::createLLVMCall(OpBuilder &b, Location loc,
+                                                   LLVM::LLVMFuncOp fn,
+                                                   ValueRange paramTypes,
+                                                   ArrayRef<Type> resultTypes) {
+  return b
+      .create<LLVM::CallOp>(loc, resultTypes, b.getSymbolRefAttr(fn),
+                            paramTypes)
+      ->getResults();
+}