Ensure CallOp with SPIR_FUNC cconv (#3692)

This PR ensures `CallOp` has `SPIR_FUNC` cconv on SPIR target.
By adding a pattern to fix-up calling convention, more common patterns
can be reused.
This PR also fixes cases of `CallOp` without `SPIR_FUNC` calling
convention.

---------

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

test/Conversion/intel/tritongpu_to_gen.mlir

@@ -1985,8 +1985,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // CHECK-LABEL: print_ptr
   tt.func @print_ptr(%arg0 : tensor<256x!tt.ptr<i32>, #blocked0>) {
-    // CHECK: llvm.call @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32, !llvm.ptr<1>) -> i32
-    // CHECK-NEXT: llvm.call @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32, !llvm.ptr<1>) -> i32
+    // CHECK: llvm.call spir_funccc @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32, !llvm.ptr<1>) -> i32
+    // CHECK-NEXT: llvm.call spir_funccc @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32, !llvm.ptr<1>) -> i32
     tt.print "ptr: " {hex = false, isSigned = array<i32: 0>} : %arg0 : tensor<256x!tt.ptr<i32>, #blocked0>
     tt.return
   }
@@ -1998,7 +1998,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // Test that %u format specifier is used if isSigned is false
   // CHECK: llvm.mlir.global internal constant @printfFormat_("pid (%u, %u, %u) idx ()int32 tensor: %u\0A\00") {addr_space = 2 : i32}
   // CHECK-LABEL: print_int32_tensor_issigned_off
-  // CHECK: llvm.call @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32) -> i32
+  // CHECK: llvm.call spir_funccc @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32) -> i32
   tt.func @print_int32_tensor_issigned_off(%arg0 : i32) {
     tt.print "int32 tensor: " {hex = false, isSigned = array<i32: 0>} : %arg0 : i32
     tt.return
@@ -2011,7 +2011,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   // Test that %i format specifier is used if isSigned is true
   // CHECK: llvm.mlir.global internal constant @printfFormat_("pid (%u, %u, %u) idx ()int32 tensor: %i\0A\00") {addr_space = 2 : i32}
   // CHECK-LABEL: print_int32_tensor_issigned_on
-  // CHECK: llvm.call @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32) -> i32
+  // CHECK: llvm.call spir_funccc @_Z18__spirv_ocl_printf(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) vararg(!llvm.func<i32 (ptr<2>, ...)>) : (!llvm.ptr<2>, i32, i32, i32, i32) -> i32
   tt.func @print_int32_tensor_issigned_on(%arg0 : i32) {
     tt.print "int32 tensor: " {hex = false, isSigned = array<i32: 1>} : %arg0 : i32
     tt.return

test/TritonIntelGPU/tritonintelgpu-rewrite-stack-ptr.mlir

@@ -7,7 +7,7 @@ module attributes {triton_intel_gpu.min_sg_size = 16 : i32, triton_intel_gpu.sup
     %0 = tt.load %arg0 : !tt.ptr<f32>
     %1 = tt.load %arg1 : !tt.ptr<f32>
     // CHECK: llvm.mlir.poison : !llvm.ptr<3>
-    // CHECK: llvm.call @noinline_simple_fn__fp32_fp32_Pfp32__(%8, %17, %arg2, %18, %arg2)
+    // CHECK: llvm.call spir_funccc @noinline_simple_fn__fp32_fp32_Pfp32__(%8, %17, %arg2, %18, %arg2)
     tt.call @noinline_simple_fn__fp32_fp32_Pfp32__(%0, %1, %arg2) : (f32, f32, !tt.ptr<f32>) -> ()
     tt.return
   }
@@ -31,7 +31,7 @@ module attributes {triton_intel_gpu.min_sg_size = 16 : i32, triton_intel_gpu.sup
   tt.func public @kernel(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}) attributes {noinline = false} {
     %0 = tt.load %arg0 : !tt.ptr<f32>
     %1 = tt.load %arg1 : !tt.ptr<f32>
-    // CHECK: llvm.call @noinline_shared_fn__fp32_fp32_Pfp32__(%8, %17, %arg2, %arg3, %arg2)
+    // CHECK: llvm.call spir_funccc @noinline_shared_fn__fp32_fp32_Pfp32__(%8, %17, %arg2, %arg3, %arg2)
     tt.call @noinline_shared_fn__fp32_fp32_Pfp32__(%0, %1, %arg2) {allocation.offset = 0 : i32} : (f32, f32, !tt.ptr<f32>) -> ()
     tt.return
   }

test/TritonIntelGPU/tritonintlgpu-nested-layout.mlir

@@ -57,6 +57,12 @@ inline bool hasSpirvTargetArch(Operation *op) {
       triton::gpu::intel::TritonIntelGPUDialect::getTargetArchAttrName());
   return !arch || arch.str().substr(0, 4) == "spir";
 }
+
+inline LLVM::cconv::CConv getRequiredCConv(Operation *op) {
+  if (hasSpirvTargetArch(op))
+    return LLVM::cconv::CConv::SPIR_FUNC;
+  llvm_unreachable("Unexpected target architecture");
+}
 } // namespace mlir::triton::gpu::intel
 
 #endif // TRITON_DIALECT_TRITON_INTEL_GPU_IR_UTILS_H

third_party/intel/include/Dialect/TritonIntelGPU/IR/Utils.h

@@ -1,6 +1,7 @@
 add_triton_library(TritonIntelGPUToLLVM
     ArithOpsToLLVM.cpp
     BF16Casts.cpp
+    ControlFlowOpToLLVM.cpp
     ConvertLayoutOpToLLVM.cpp
     DecomposeUnsupportedConversions.cpp
     DotOpToLLVM/DPAS.cpp

third_party/intel/lib/TritonIntelGPUToLLVM/CMakeLists.txt

@@ -0,0 +1,27 @@
+#include "PatternTritonGPUOpToLLVM.h"
+#include "intel/include/Dialect/TritonIntelGPU/IR/Utils.h"
+
+namespace {
+
+struct FixCallCConv : public ConvertOpToLLVMPattern<LLVM::CallOp> {
+  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(LLVM::CallOp op, LLVM::CallOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.startOpModification(op);
+    op.setCConv(triton::gpu::intel::getRequiredCConv(op));
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::triton::intel::populateControlFlowOpToLLVMPattern(
+    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+    const TargetInfoBase &targetInfo, PatternBenefit benefit) {
+  patterns.add<FixCallCConv>(typeConverter);
+  mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
+                                                   targetInfo, benefit);
+}

third_party/intel/lib/TritonIntelGPUToLLVM/ControlFlowOpToLLVM.cpp

@@ -1085,34 +1085,6 @@ struct FpToFpOpConversion
   }
 };
 
-struct ExternElementwiseOpConversion
-    : public ElementwiseOpConversionBase<ExternElementwiseOp,
-                                         ExternElementwiseOpConversion> {
-  using Base = ElementwiseOpConversionBase<ExternElementwiseOp,
-                                           ExternElementwiseOpConversion>;
-  using Base::Base;
-  using Adaptor = typename Base::OpAdaptor;
-  typedef typename Base::OpAdaptor OpAdaptor;
-
-  SmallVector<Value> createDestOps(ExternElementwiseOp op, OpAdaptor adaptor,
-                                   ConversionPatternRewriter &rewriter,
-                                   Type elemTy, MultipleOperandsRange operands,
-                                   Location loc) const {
-    StringRef funcName = op.getSymbol();
-    if (funcName.empty())
-      llvm::errs() << "ExternElementwiseOpConversion";
-
-    Type funcType = getFunctionType(elemTy, operands[0]);
-    LLVM::LLVMFuncOp funcOp = appendOrGetExternFuncOp(
-        rewriter, op, funcName, funcType, op.getLibname(), op.getLibpath());
-
-    auto callOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, operands[0]);
-    callOp.setCConv(LLVM::cconv::CConv::SPIR_FUNC);
-
-    return {callOp.getResult()};
-  }
-};
-
 template <typename SourceOp, typename DestOp>
 struct ElementwiseOpConversion
     : ElementwiseOpConversionBase<SourceOp,
@@ -1292,38 +1264,6 @@ struct AbsFOpConversion
   }
 };
 
-struct MulhiUIOpConversion
-    : public ElementwiseOpConversionBase<MulhiUIOp, MulhiUIOpConversion> {
-  using Base = ElementwiseOpConversionBase<MulhiUIOp, MulhiUIOpConversion>;
-  using Base::Base;
-  using Adaptor = typename Base::OpAdaptor;
-  explicit MulhiUIOpConversion(LLVMTypeConverter &typeConverter,
-                               ModuleAxisInfoAnalysis &axisAnalysisPass,
-                               const TargetInfoBase &targetInfo,
-                               PatternBenefit benefit = 1)
-      : ElementwiseOpConversionBase(typeConverter, axisAnalysisPass, benefit),
-        targetInfo(targetInfo) {}
-  SmallVector<Value> createDestOps(MulhiUIOp op, Adaptor adaptor,
-                                   ConversionPatternRewriter &rewriter,
-                                   Type elemTy, MultipleOperandsRange operands,
-                                   Location loc) const {
-
-    Type resultElementTy = getElementTypeOrSelf(op.getResult().getType());
-    assert(resultElementTy.isInteger(32) || resultElementTy.isInteger(64));
-
-    std::string funcName = targetInfo.getMulhiFuncName(resultElementTy);
-    Type funcType = getFunctionType(elemTy, operands[0]);
-    LLVM::LLVMFuncOp funcOp =
-        appendOrGetExternFuncOp(rewriter, op, funcName, funcType);
-    auto callOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, operands[0]);
-    callOp.setCConv(LLVM::cconv::CConv::SPIR_FUNC);
-    return {callOp.getResult()};
-  }
-
-protected:
-  const TargetInfoBase &targetInfo;
-};
-
 struct PreciseSqrtOpConversion
     : ElementwiseOpConversionBase<PreciseSqrtOp, PreciseSqrtOpConversion> {
   using Base =
@@ -1401,10 +1341,6 @@ void populateElementwiseOpToLLVMPatterns(
 
   mlir::triton::populateElementwiseOpToLLVMPatterns(
       typeConverter, patterns, axisInfoAnalysis, targetInfo, benefit);
-  patterns.add<MulhiUIOpConversion>(typeConverter, axisInfoAnalysis, targetInfo,
-                                    benefit);
-  patterns.add<ExternElementwiseOpConversion>(typeConverter, axisInfoAnalysis,
-                                              benefit);
 
   patterns.add<AbsFOpConversion>(typeConverter, axisInfoAnalysis, benefit);
   patterns.add<ElementwiseOpConversion<arith::DivFOp, LLVM::FDivOp>>(

third_party/intel/lib/TritonIntelGPUToLLVM/ElementwiseOpToLLVM.cpp

@@ -49,6 +49,11 @@ void populateConvertLayoutOpToLLVMPatterns(LLVMTypeConverter &typeConverter,
                                            RewritePatternSet &patterns,
                                            PatternBenefit benefit);
 
+void populateControlFlowOpToLLVMPattern(LLVMTypeConverter &typeConverter,
+                                        RewritePatternSet &patterns,
+                                        const TargetInfoBase &targetInfo,
+                                        PatternBenefit benefit);
+
 void populateSPMDOpToLLVMPattern(LLVMTypeConverter &typeConverter,
                                  RewritePatternSet &patterns,
                                  const TargetInfoBase &targetInfo,

third_party/intel/lib/TritonIntelGPUToLLVM/PatternTritonGPUOpToLLVM.h

@@ -240,8 +240,8 @@ class TritonGPUToLLVMPipelineManager {
     if (isAdvancedPathEnabled) {
       intel::populateArithOpsToLLVMPatterns(typeConverter, patterns, benefit);
       intel::populateBF16CastsLLVMPatterns(typeConverter, patterns, benefit);
-      mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
-                                                       targetInfo, benefit);
+      intel::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
+                                                targetInfo, benefit);
       intel::populateTritonOpsToLLVMPatterns(typeConverter, patterns, benefit);
     } else {
       intel::populateConvertLayoutOpToLLVMPatterns(typeConverter, targetInfo,
@@ -272,8 +272,8 @@ class TritonGPUToLLVMPipelineManager {
                                     benefit);
       mlir::triton::populateMemoryOpToLLVMPatterns(typeConverter, targetInfo,
                                                    patterns, benefit);
-      mlir::triton::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
-                                                       targetInfo, benefit);
+      intel::populateControlFlowOpToLLVMPattern(typeConverter, patterns,
+                                                targetInfo, benefit);
       mlir::triton::populateMakeRangeOpToLLVMPattern(typeConverter, targetInfo,
                                                      patterns, benefit);
       intel::populateFp4ToFpToLLVMPatterns(typeConverter, patterns, benefit);

third_party/intel/lib/TritonIntelGPUToLLVM/PipelineManager.h

@@ -221,7 +221,8 @@ void TargetInfo::printf(RewriterBase &rewriter, Value formatStrStart,
     operands.push_back(printfPromoteValue(
         rewriter, arg, isSigned.empty() ? true : isSigned[i]));
   }
-  b.call(funcOp, operands);
+  auto callOp = b.call(funcOp, operands);
+  callOp.setCConv(triton::gpu::intel::getRequiredCConv(callOp));
 }
 
 void TargetInfo::printf(RewriterBase &rewriter, StringRef msg, ValueRange args,