Use puts for cf.assert and printString for vector.print

Author: MacDue

mlir/include/mlir/Conversion/LLVMCommon/PrintCallHelper.h

@@ -11,6 +11,7 @@
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "llvm/ADT/StringRef.h"
+#include <optional>
 
 namespace mlir {
 
@@ -20,10 +21,13 @@ class LLVMTypeConverter;
 namespace LLVM {
 
 /// Generate IR that prints the given string to stdout.
+/// If a custom runtime function is defined via `runtimeFunctionName`, it must
+/// have the signature void(char const*). The default function is `printString`.
 void createPrintStrCall(OpBuilder &builder, Location loc, ModuleOp moduleOp,
                         StringRef symbolName, StringRef string,
                         const LLVMTypeConverter &typeConverter,
-                        bool addNewline = true);
+                        bool addNewline = true,
+                        std::optional<StringRef> runtimeFunctionName = {});
 } // namespace LLVM
 
 } // namespace mlir

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

@@ -16,6 +16,7 @@
 
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
+#include <optional>
 
 namespace mlir {
 class Location;
@@ -38,8 +39,12 @@ LLVM::LLVMFuncOp lookupOrCreatePrintF16Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintBF16Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintF32Fn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintF64Fn(ModuleOp moduleOp);
-LLVM::LLVMFuncOp lookupOrCreatePrintStringFn(ModuleOp moduleOp,
-                                             bool opaquePointers);
+/// Declares a function to print a C-string.
+/// If a custom runtime function is defined via `runtimeFunctionName`, it must
+/// have the signature void(char const*). The default function is `printString`.
+LLVM::LLVMFuncOp
+lookupOrCreatePrintStringFn(ModuleOp moduleOp, bool opaquePointers,
+                            std::optional<StringRef> runtimeFunctionName = {});
 LLVM::LLVMFuncOp lookupOrCreatePrintOpenFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCloseFn(ModuleOp moduleOp);
 LLVM::LLVMFuncOp lookupOrCreatePrintCommaFn(ModuleOp moduleOp);

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

@@ -62,7 +62,8 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
     // Failed block: Generate IR to print the message and call `abort`.
     Block *failureBlock = rewriter.createBlock(opBlock->getParent());
     LLVM::createPrintStrCall(rewriter, loc, module, "assert_msg", op.getMsg(),
-                             *getTypeConverter());
+                             *getTypeConverter(), /*addNewLine=*/false,
+                             /*runtimeFunctionName=*/"puts");
     if (abortOnFailedAssert) {
       // Insert the `abort` declaration if necessary.
       auto abortFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("abort");

mlir/lib/Conversion/LLVMCommon/PrintCallHelper.cpp

@@ -27,11 +27,10 @@ static std::string ensureSymbolNameIsUnique(ModuleOp moduleOp,
   return uniqueName;
 }
 
-void mlir::LLVM::createPrintStrCall(OpBuilder &builder, Location loc,
-                                    ModuleOp moduleOp, StringRef symbolName,
-                                    StringRef string,
-                                    const LLVMTypeConverter &typeConverter,
-                                    bool addNewline) {
+void mlir::LLVM::createPrintStrCall(
+    OpBuilder &builder, Location loc, ModuleOp moduleOp, StringRef symbolName,
+    StringRef string, const LLVMTypeConverter &typeConverter, bool addNewline,
+    std::optional<StringRef> runtimeFunctionName) {
   auto ip = builder.saveInsertionPoint();
   builder.setInsertionPointToStart(moduleOp.getBody());
   MLIRContext *ctx = builder.getContext();
@@ -61,7 +60,7 @@ void mlir::LLVM::createPrintStrCall(OpBuilder &builder, Location loc,
       loc, typeConverter.getPointerType(builder.getI8Type()), arrayTy, msgAddr,
       indices);
   Operation *printer = LLVM::lookupOrCreatePrintStringFn(
-      moduleOp, typeConverter.useOpaquePointers());
+      moduleOp, typeConverter.useOpaquePointers(), runtimeFunctionName);
   builder.create<LLVM::CallOp>(loc, TypeRange(), SymbolRefAttr::get(printer),
                                gep);
 }

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

@@ -107,9 +107,10 @@ static LLVM::LLVMPointerType getVoidPtr(MLIRContext *context,
   return getCharPtr(context, opaquePointers);
 }
 
-LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintStringFn(ModuleOp moduleOp,
-                                                         bool opaquePointers) {
-  return lookupOrCreateFn(moduleOp, kPrintString,
+LLVM::LLVMFuncOp mlir::LLVM::lookupOrCreatePrintStringFn(
+    ModuleOp moduleOp, bool opaquePointers,
+    std::optional<StringRef> runtimeFunctionName) {
+  return lookupOrCreateFn(moduleOp, runtimeFunctionName.value_or(kPrintString),
                           getCharPtr(moduleOp->getContext(), opaquePointers),
                           LLVM::LLVMVoidType::get(moduleOp->getContext()));
 }

mlir/lib/ExecutionEngine/RunnerUtils.cpp

@@ -158,10 +158,8 @@ extern "C" void printMemrefC64(int64_t rank, void *ptr) {
   _mlir_ciface_printMemrefC64(&descriptor);
 }
 
+/// Deprecated. This should be unified with printString from CRunnerUtils.
 extern "C" void printCString(char *str) { fputs(str, stdout); }
-extern "C" void printString(char const *str) {
-  printCString(const_cast<char *>(str));
-}
 
 extern "C" void _mlir_ciface_printMemref0dF32(StridedMemRefType<float, 0> *M) {
   impl::printMemRef(*M);

mlir/test/Conversion/ControlFlowToLLVM/assert.mlir

@@ -10,12 +10,12 @@ func.func @main() {
   return
 }
 
-// CHECK: llvm.func @printString(!llvm.ptr)
+// CHECK: llvm.func @puts(!llvm.ptr)
 
 // CHECK-LABEL: @main
 // CHECK: llvm.cond_br %{{.*}}, ^{{.*}}, ^[[FALSE_BRANCH:[[:alnum:]]+]]
 
 // CHECK: ^[[FALSE_BRANCH]]:
 // CHECK: %[[ADDRESS_OF:.*]] = llvm.mlir.addressof @{{.*}} : !llvm.ptr{{$}}
 // CHECK: %[[GEP:.*]] = llvm.getelementptr %[[ADDRESS_OF]][0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<{{[0-9]+}} x i8>
-// CHECK: llvm.call @printString(%[[GEP]]) : (!llvm.ptr) -> ()
+// CHECK: llvm.call @puts(%[[GEP]]) : (!llvm.ptr) -> ()

mlir/test/Integration/Dialect/ControlFlow/assert.mlir

@@ -1,6 +1,6 @@
 // RUN: mlir-opt %s -test-cf-assert \
 // RUN:     -convert-func-to-llvm | \
-// RUN: mlir-cpu-runner -e main -entry-point-result=void -shared-libs=%mlir_runner_utils | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void | \
 // RUN: FileCheck %s
 
 func.func @main() {