[MLIR][GPU] Define gpu.printf op and its lowerings

- Define a gpu.printf op, which can be lowered to any GPU printf() support (which is present in CUDA, HIP, and OpenCL). This op only supports constant format strings and scalar arguments
- Define the lowering of gpu.pirntf to a call to printf() (which is what is required for AMD GPUs when using OpenCL) as well as to the hostcall interface present in the AMD Open Compute device library, which is the interface present when kernels are running under HIP.
- Add a "runtime" enum that allows specifying which of the possible runtimes a ROCDL kernel will be executed under or that the runtime is unknown. This enum controls how gpu.printf is lowered

This change does not enable lowering for Nvidia GPUs, but such a lowering should be possible in principle.

And:
[MLIR][AMDGPU] Always set amdgpu-implicitarg-num-bytes=56 on kernels

This is something that Clang always sets on both OpenCL and HIP kernels, and failing to include it causes mysterious crashes with printf() support.

In addition, revert the max-flat-work-group-size to (1, 256) to avoid triggering bugs in the AMDGPU backend.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D110448

mlir/include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h

@@ -8,6 +8,7 @@
 #ifndef MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
 #define MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
 
+#include "mlir/Conversion/GPUToROCDL/Runtimes.h"
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
 #include <memory>
 
@@ -25,8 +26,11 @@ class GPUModuleOp;
 } // namespace gpu
 
 /// Collect a set of patterns to convert from the GPU dialect to ROCDL.
+/// If `runtime` is Unknown, gpu.printf will not be lowered
+/// The resulting pattern set should be run over a gpu.module op
 void populateGpuToROCDLConversionPatterns(LLVMTypeConverter &converter,
-                                          RewritePatternSet &patterns);
+                                          RewritePatternSet &patterns,
+                                          gpu::amd::Runtime runtime);
 
 /// Configure target to convert from the GPU dialect to ROCDL.
 void configureGpuToROCDLConversionLegality(ConversionTarget &target);
@@ -36,7 +40,8 @@ void configureGpuToROCDLConversionLegality(ConversionTarget &target);
 /// is configurable.
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
 createLowerGpuOpsToROCDLOpsPass(
-    unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout);
+    unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout,
+    gpu::amd::Runtime runtime = gpu::amd::Runtime::Unknown);
 
 } // namespace mlir
 

mlir/include/mlir/Conversion/GPUToROCDL/Runtimes.h

@@ -0,0 +1,24 @@
+//===- Runtimes.h - Possible runtimes for AMD GPUs ---*- 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
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_GPUTOROCDL_RUNTIMES_H
+#define MLIR_CONVERSION_GPUTOROCDL_RUNTIMES_H
+
+namespace mlir {
+namespace gpu {
+namespace amd {
+/// Potential runtimes for AMD GPU kernels
+enum Runtime {
+  Unknown = 0,
+  HIP = 1,
+  OpenCL = 2,
+};
+} // end namespace amd
+} // end namespace gpu
+} // end namespace mlir
+
+#endif // MLIR_CONVERSION_GPUTOROCDL_RUNTIMES_H

mlir/include/mlir/Conversion/Passes.td

@@ -203,7 +203,15 @@ def ConvertGpuOpsToROCDLOps : Pass<"convert-gpu-to-rocdl", "gpu::GPUModuleOp"> {
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
-           "Bitwidth of the index type, 0 to use size of machine word">
+           "Bitwidth of the index type, 0 to use size of machine word">,
+    Option<"runtime", "runtime", "::mlir::gpu::amd::Runtime",
+          "::mlir::gpu::amd::Runtime::Unknown",
+          "Runtime code will be run on (default is Unknown, can also use HIP or OpenCl)",
+          [{::llvm::cl::values(
+            clEnumValN(::mlir::gpu::amd::Runtime::Unknown, "unknown", "Unknown (default)"),
+            clEnumValN(::mlir::gpu::amd::Runtime::HIP, "HIP", "HIP"),
+            clEnumValN(::mlir::gpu::amd::Runtime::OpenCL, "OpenCL", "OpenCL")
+          )}]>
   ];
 }
 

mlir/include/mlir/Dialect/GPU/GPUOps.td

@@ -547,6 +547,22 @@ def GPU_LaunchOp : GPU_Op<"launch">,
   let hasCanonicalizer = 1;
 }
 
+def GPU_PrintfOp : GPU_Op<"printf", [MemoryEffects<[MemWrite]>]>,
+  Arguments<(ins StrAttr:$format,
+                Variadic<AnyTypeOf<[AnyInteger, Index, AnyFloat]>>:$args)> {
+  let summary = "Device-side printf, as in CUDA or OpenCL, for debugging";
+  let description = [{
+    `gpu.printf` takes a literal format string `format` and an arbitrary number of
+    scalar arguments that should be printed.
+
+    The format string is a C-style printf string, subject to any restrictions
+    imposed by one's target platform.
+  }];
+  let assemblyFormat = [{
+    $format attr-dict ($args^ `:` type($args))?
+  }];
+}
+
 def GPU_ReturnOp : GPU_Op<"return", [HasParent<"GPUFuncOp">, NoSideEffect,
                                      Terminator]>,
     Arguments<(ins Variadic<AnyType>:$operands)>, Results<(outs)> {

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "GPUOpsLowering.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -144,3 +145,200 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
   rewriter.eraseOp(gpuFuncOp);
   return success();
 }
+
+static const char formatStringPrefix[] = "printfFormat_";
+
+template <typename T>
+static LLVM::LLVMFuncOp getOrDefineFunction(T &moduleOp, const Location loc,
+                                            ConversionPatternRewriter &rewriter,
+                                            StringRef name,
+                                            LLVM::LLVMFunctionType type) {
+  LLVM::LLVMFuncOp ret;
+  if (!(ret = moduleOp.template lookupSymbol<LLVM::LLVMFuncOp>(name))) {
+    ConversionPatternRewriter::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(moduleOp.getBody());
+    ret = rewriter.create<LLVM::LLVMFuncOp>(loc, name, type,
+                                            LLVM::Linkage::External);
+  }
+  return ret;
+}
+
+LogicalResult GPUPrintfOpToHIPLowering::matchAndRewrite(
+    gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  Location loc = gpuPrintfOp->getLoc();
+
+  mlir::Type llvmI8 = typeConverter->convertType(rewriter.getI8Type());
+  mlir::Type i8Ptr = LLVM::LLVMPointerType::get(llvmI8);
+  mlir::Type llvmIndex = typeConverter->convertType(rewriter.getIndexType());
+  mlir::Type llvmI32 = typeConverter->convertType(rewriter.getI32Type());
+  mlir::Type llvmI64 = typeConverter->convertType(rewriter.getI64Type());
+  // Note: this is the GPUModule op, not the ModuleOp that surrounds it
+  // This ensures that global constants and declarations are placed within
+  // the device code, not the host code
+  auto moduleOp = gpuPrintfOp->getParentOfType<gpu::GPUModuleOp>();
+
+  auto ocklBegin =
+      getOrDefineFunction(moduleOp, loc, rewriter, "__ockl_printf_begin",
+                          LLVM::LLVMFunctionType::get(llvmI64, {llvmI64}));
+  LLVM::LLVMFuncOp ocklAppendArgs;
+  if (!adaptor.args().empty()) {
+    ocklAppendArgs = getOrDefineFunction(
+        moduleOp, loc, rewriter, "__ockl_printf_append_args",
+        LLVM::LLVMFunctionType::get(
+            llvmI64, {llvmI64, /*numArgs*/ llvmI32, llvmI64, llvmI64, llvmI64,
+                      llvmI64, llvmI64, llvmI64, llvmI64, /*isLast*/ llvmI32}));
+  }
+  auto ocklAppendStringN = getOrDefineFunction(
+      moduleOp, loc, rewriter, "__ockl_printf_append_string_n",
+      LLVM::LLVMFunctionType::get(
+          llvmI64,
+          {llvmI64, i8Ptr, /*length (bytes)*/ llvmI64, /*isLast*/ llvmI32}));
+
+  /// Start the printf hostcall
+  Value zeroI64 = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmI64, rewriter.getI64IntegerAttr(0));
+  auto printfBeginCall = rewriter.create<LLVM::CallOp>(loc, ocklBegin, zeroI64);
+  Value printfDesc = printfBeginCall.getResult(0);
+
+  // Create a global constant for the format string
+  unsigned stringNumber = 0;
+  SmallString<16> stringConstName;
+  do {
+    stringConstName.clear();
+    (formatStringPrefix + Twine(stringNumber++)).toStringRef(stringConstName);
+  } while (moduleOp.lookupSymbol(stringConstName));
+
+  llvm::SmallString<20> formatString(adaptor.format().getValue());
+  formatString.push_back('\0'); // Null terminate for C
+  size_t formatStringSize = formatString.size_in_bytes();
+
+  auto globalType = LLVM::LLVMArrayType::get(llvmI8, formatStringSize);
+  LLVM::GlobalOp global;
+  {
+    ConversionPatternRewriter::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(moduleOp.getBody());
+    global = rewriter.create<LLVM::GlobalOp>(
+        loc, globalType,
+        /*isConstant=*/true, LLVM::Linkage::Internal, stringConstName,
+        rewriter.getStringAttr(formatString));
+  }
+
+  // Get a pointer to the format string's first element and pass it to printf()
+  Value globalPtr = rewriter.create<LLVM::AddressOfOp>(loc, global);
+  Value zero = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmIndex, rewriter.getIntegerAttr(llvmIndex, 0));
+  Value stringStart = rewriter.create<LLVM::GEPOp>(
+      loc, i8Ptr, globalPtr, mlir::ValueRange({zero, zero}));
+  Value stringLen = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmI64, rewriter.getI64IntegerAttr(formatStringSize));
+
+  Value oneI32 = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmI32, rewriter.getI32IntegerAttr(1));
+  Value zeroI32 = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmI32, rewriter.getI32IntegerAttr(0));
+
+  mlir::ValueRange appendFormatArgs = {printfDesc, stringStart, stringLen,
+                                       adaptor.args().empty() ? oneI32
+                                                              : zeroI32};
+  auto appendFormatCall =
+      rewriter.create<LLVM::CallOp>(loc, ocklAppendStringN, appendFormatArgs);
+  printfDesc = appendFormatCall.getResult(0);
+
+  // __ockl_printf_append_args takes 7 values per append call
+  constexpr size_t argsPerAppend = 7;
+  size_t nArgs = adaptor.args().size();
+  for (size_t group = 0; group < nArgs; group += argsPerAppend) {
+    size_t bound = std::min(group + argsPerAppend, nArgs);
+    size_t numArgsThisCall = bound - group;
+
+    SmallVector<mlir::Value, 2 + argsPerAppend + 1> arguments;
+    arguments.push_back(printfDesc);
+    arguments.push_back(rewriter.create<LLVM::ConstantOp>(
+        loc, llvmI32, rewriter.getI32IntegerAttr(numArgsThisCall)));
+    for (size_t i = group; i < bound; ++i) {
+      Value arg = adaptor.args()[i];
+      if (auto floatType = arg.getType().dyn_cast<FloatType>()) {
+        if (!floatType.isF64())
+          arg = rewriter.create<LLVM::FPExtOp>(
+              loc, typeConverter->convertType(rewriter.getF64Type()), arg);
+        arg = rewriter.create<LLVM::BitcastOp>(loc, llvmI64, arg);
+      }
+      if (arg.getType().getIntOrFloatBitWidth() != 64)
+        arg = rewriter.create<LLVM::ZExtOp>(loc, llvmI64, arg);
+
+      arguments.push_back(arg);
+    }
+    // Pad out to 7 arguments since the hostcall always needs 7
+    for (size_t extra = numArgsThisCall; extra < argsPerAppend; ++extra) {
+      arguments.push_back(zeroI64);
+    }
+
+    auto isLast = (bound == nArgs) ? oneI32 : zeroI32;
+    arguments.push_back(isLast);
+    auto call = rewriter.create<LLVM::CallOp>(loc, ocklAppendArgs, arguments);
+    printfDesc = call.getResult(0);
+  }
+  rewriter.eraseOp(gpuPrintfOp);
+  return success();
+}
+
+LogicalResult GPUPrintfOpToLLVMCallLowering::matchAndRewrite(
+    gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  Location loc = gpuPrintfOp->getLoc();
+
+  mlir::Type llvmI8 = typeConverter->convertType(rewriter.getIntegerType(8));
+  mlir::Type i8Ptr = LLVM::LLVMPointerType::get(llvmI8, addressSpace);
+  mlir::Type llvmIndex = typeConverter->convertType(rewriter.getIndexType());
+
+  // Note: this is the GPUModule op, not the ModuleOp that surrounds it
+  // This ensures that global constants and declarations are placed within
+  // the device code, not the host code
+  auto moduleOp = gpuPrintfOp->getParentOfType<gpu::GPUModuleOp>();
+
+  auto printfType = LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {i8Ptr},
+                                                /*isVarArg=*/true);
+  LLVM::LLVMFuncOp printfDecl =
+      getOrDefineFunction(moduleOp, loc, rewriter, "printf", printfType);
+
+  // Create a global constant for the format string
+  unsigned stringNumber = 0;
+  SmallString<16> stringConstName;
+  do {
+    stringConstName.clear();
+    (formatStringPrefix + Twine(stringNumber++)).toStringRef(stringConstName);
+  } while (moduleOp.lookupSymbol(stringConstName));
+
+  llvm::SmallString<20> formatString(adaptor.format().getValue());
+  formatString.push_back('\0'); // Null terminate for C
+  auto globalType =
+      LLVM::LLVMArrayType::get(llvmI8, formatString.size_in_bytes());
+  LLVM::GlobalOp global;
+  {
+    ConversionPatternRewriter::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPointToStart(moduleOp.getBody());
+    global = rewriter.create<LLVM::GlobalOp>(
+        loc, globalType,
+        /*isConstant=*/true, LLVM::Linkage::Internal, stringConstName,
+        rewriter.getStringAttr(formatString), /*allignment=*/0, addressSpace);
+  }
+
+  // Get a pointer to the format string's first element
+  Value globalPtr = rewriter.create<LLVM::AddressOfOp>(loc, global);
+  Value zero = rewriter.create<LLVM::ConstantOp>(
+      loc, llvmIndex, rewriter.getIntegerAttr(llvmIndex, 0));
+  Value stringStart = rewriter.create<LLVM::GEPOp>(
+      loc, i8Ptr, globalPtr, mlir::ValueRange({zero, zero}));
+
+  // Construct arguments and function call
+  auto argsRange = adaptor.args();
+  SmallVector<Value, 4> printfArgs;
+  printfArgs.reserve(argsRange.size() + 1);
+  printfArgs.push_back(stringStart);
+  printfArgs.append(argsRange.begin(), argsRange.end());
+
+  rewriter.create<LLVM::CallOp>(loc, printfDecl, printfArgs);
+  rewriter.eraseOp(gpuPrintfOp);
+  return success();
+}

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h

@@ -33,6 +33,40 @@ struct GPUFuncOpLowering : ConvertOpToLLVMPattern<gpu::GPUFuncOp> {
   StringAttr kernelAttributeName;
 };
 
+/// The lowering of gpu.printf to a call to HIP hostcalls
+///
+/// Simplifies llvm/lib/Transforms/Utils/AMDGPUEmitPrintf.cpp, as we don't have
+/// to deal with %s (even if there were first-class strings in MLIR, they're not
+/// legal input to gpu.printf) or non-constant format strings
+struct GPUPrintfOpToHIPLowering : public ConvertOpToLLVMPattern<gpu::PrintfOp> {
+  using ConvertOpToLLVMPattern<gpu::PrintfOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
+/// The lowering of gpu.printf to a call to an external printf() function
+///
+/// This pass will add a declaration of printf() to the GPUModule if needed
+/// and seperate out the format strings into global constants. For some
+/// runtimes, such as OpenCL on AMD, this is sufficient setup, as the compiler
+/// will lower printf calls to appropriate device-side code
+struct GPUPrintfOpToLLVMCallLowering
+    : public ConvertOpToLLVMPattern<gpu::PrintfOp> {
+  GPUPrintfOpToLLVMCallLowering(LLVMTypeConverter &converter,
+                                int addressSpace = 0)
+      : ConvertOpToLLVMPattern<gpu::PrintfOp>(converter),
+        addressSpace(addressSpace) {}
+
+  LogicalResult
+  matchAndRewrite(gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+
+private:
+  int addressSpace;
+};
+
 struct GPUReturnOpLowering : public ConvertOpToLLVMPattern<gpu::ReturnOp> {
   using ConvertOpToLLVMPattern<gpu::ReturnOp>::ConvertOpToLLVMPattern;