[mlir] Introduce a pattern to lower `gpu.subgroup_reduce` to `nvvm.redux_op`

[guraypp]

Summary

This revision introduces a pattern to lower gpu.subgroup_reduce op into to the nvvm.redux_sync op. The op must be run by the entire subgroup, otherwise it is undefined behaviour.

It also adds a flag and populate function, because the op is not avaiable for every gpu (sm80+), so it can be used when it is desired.

Depends on D142088

[Hardcode84]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

76

Instead of hard crashing should probably return optional and check it in pattern.

92

missing return

[guraypp]

address comments

[guraypp]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

76

good catch, thanks!

[nicolasvasilache]

Looks good once we purge errors

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

92

not an error but rewriter.notifyMatchFailure everywhere plz

[guraypp]

Change emitErrors to notifyMatchFailure

mlir/include/mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h

 
 /// Configure target to convert from the GPU dialect to NVVM.
 void configureGpuToNVVMConversionLegality(ConversionTarget &target);
 
 /// Collect a set of patterns to convert from the GPU dialect to NVVM.
 void populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                          RewritePatternSet &patterns);
 
+/// Populate GpuSubgroupReduce pattern to NVVM. It generates a specific nvvm
+/// op that is not available on every GPU.
+void populateGpuSubgroupReduceOpLoweringPattern(LLVMTypeConverter &converter,
+                                                RewritePatternSet &patterns);
+
 /// Collect a set of patterns to convert WMMA ops from GPU dialect to NVVM.
 void populateGpuWMMAToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                              RewritePatternSet &patterns);
 
 /// Creates a pass that lowers GPU dialect operations to NVVM counterparts. The
 /// index bitwidth used for the lowering of the device side index computations
 /// is configurable.
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> createLowerGpuOpsToNVVMOpsPass(
-    unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout);
+    unsigned indexBitwidth = kDeriveIndexBitwidthFromDataLayout,
+    bool hasRedux = false);
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUTONVVM_GPUTONVVMPASS_H_

mlir/include/mlir/Conversion/Passes.td

def ConvertGpuOpsToNVVMOps : Pass<"convert-gpu-to-nvvm", "gpu::GPUModuleOp"> {
   let dependentDialects = [
     "cf::ControlFlowDialect",
     "memref::MemRefDialect",
     "NVVM::NVVMDialect",
   ];
   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<"hasRedux", "has-redux", "bool", /*default=*/"false",
+           "Target gpu supports redux">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToROCDL
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuOpsToROCDLOps : Pass<"convert-gpu-to-rocdl", "gpu::GPUModuleOp"> {

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

static NVVM::ShflKind convertShflKind(gpu::ShuffleMode mode) {
   case gpu::ShuffleMode::DOWN:
     return NVVM::ShflKind::down;
   case gpu::ShuffleMode::IDX:
     return NVVM::ShflKind::idx;
   }
   llvm_unreachable("unknown shuffle mode");
 }
 
+static Optional<NVVM::ReduxKind>
+convertReduxKind(gpu::AllReduceOperation mode) {
+  switch (mode) {
+  case gpu::AllReduceOperation::ADD:
+    return NVVM::ReduxKind::ADD;
+  case gpu::AllReduceOperation::AND:
+    return NVVM::ReduxKind::AND;
+  case gpu::AllReduceOperation::MAX:
+    return NVVM::ReduxKind::MAX;
+  case gpu::AllReduceOperation::MIN:
+    return NVVM::ReduxKind::MIN;
+  case gpu::AllReduceOperation::OR:
+    return NVVM::ReduxKind::OR;
+  case gpu::AllReduceOperation::XOR:
+    return NVVM::ReduxKind::XOR;
+  case gpu::AllReduceOperation::MUL:
+    return std::nullopt;
+  }
+  return std::nullopt;
+}
+
+/// This pass lowers gpu.subgroup_reduce op into to the nvvm.redux op. The op
+/// must be run by the entire subgroup, otherwise it is undefined behaviour.
+struct GPUSubgroupReduceOpLowering
+    : public ConvertOpToLLVMPattern<gpu::SubgroupReduceOp> {
+  using ConvertOpToLLVMPattern<gpu::SubgroupReduceOp>::ConvertOpToLLVMPattern;
+  LogicalResult
+
+  matchAndRewrite(gpu::SubgroupReduceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!op.getUniform())
+      return rewriter.notifyMatchFailure(
+          op, "cannot be lowered to redux as the op must be run "
+              "uniformly (entire subgroup).");
+    if (!op.getValue().getType().isInteger(32))
+      return rewriter.notifyMatchFailure(op, "unsupported data type");
+
+    Optional<NVVM::ReduxKind> mode = convertReduxKind(op.getOp());
+    if (!mode.has_value())
+      return rewriter.notifyMatchFailure(
+          op, "unsupported reduction mode for redux");
+
+    Location loc = op->getLoc();
+    auto int32Type = IntegerType::get(rewriter.getContext(), 32);
+    Value offset = rewriter.create<LLVM::ConstantOp>(loc, int32Type, -1);
+
+    auto reduxOp = rewriter.create<NVVM::ReduxOp>(loc, int32Type, op.getValue(),
+                                                  mode.value(), offset);
+
+    rewriter.replaceOp(op, reduxOp->getResult(0));
+    return success();
+  }
+};
+
 struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
   using ConvertOpToLLVMPattern<gpu::ShuffleOp>::ConvertOpToLLVMPattern;
 
   /// Lowers a shuffle to the corresponding NVVM op.
   ///
   /// Convert the `width` argument into an activeMask (a bitmask which specifies
   /// which threads participate in the shuffle) and a maskAndClamp (specifying
   /// the highest lane which participates in the shuffle).

 /// A pass that replaces all occurrences of GPU device operations with their
 /// corresponding NVVM equivalent.
 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
 struct LowerGpuOpsToNVVMOpsPass
     : public impl::ConvertGpuOpsToNVVMOpsBase<LowerGpuOpsToNVVMOpsPass> {
   LowerGpuOpsToNVVMOpsPass() = default;
-  LowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
+  LowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth, bool hasRedux = false) {
     this->indexBitwidth = indexBitwidth;
+    this->hasRedux = hasRedux;
   }
 
   void runOnOperation() override {
     gpu::GPUModuleOp m = getOperation();
 
     // Request C wrapper emission.
     for (auto func : m.getOps<func::FuncOp>()) {
       func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
void runOnOperation() override {
     RewritePatternSet llvmPatterns(m.getContext());
 
     arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);
     populateFuncToLLVMConversionPatterns(converter, llvmPatterns);
     populateMemRefToLLVMConversionPatterns(converter, llvmPatterns);
     populateGpuToNVVMConversionPatterns(converter, llvmPatterns);
     populateGpuWMMAToNVVMConversionPatterns(converter, llvmPatterns);
+    if (this->hasRedux)
+      populateGpuSubgroupReduceOpLoweringPattern(converter, llvmPatterns);
     LLVMConversionTarget target(getContext());
     configureGpuToNVVMConversionLegality(target);
     if (failed(applyPartialConversion(m, target, std::move(llvmPatterns))))
       signalPassFailure();
   }
 };
 
 } // namespace

 template <typename OpTy>
 static void populateOpPatterns(LLVMTypeConverter &converter,
                                RewritePatternSet &patterns, StringRef f32Func,
                                StringRef f64Func) {
   patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
   patterns.add<OpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func);
 }
 
+void mlir::populateGpuSubgroupReduceOpLoweringPattern(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
+  patterns.add<GPUSubgroupReduceOpLowering>(converter);
+}
+
 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                                RewritePatternSet &patterns) {
   populateWithGenerated(patterns);
   patterns.add<GPUPrintfOpToVPrintfLowering>(converter);
   patterns
       .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,
                                        NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>,
            GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp,
void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
   populateOpPatterns<math::SqrtOp>(converter, patterns, "__nv_sqrtf",
                                    "__nv_sqrt");
   populateOpPatterns<math::TanhOp>(converter, patterns, "__nv_tanhf",
                                    "__nv_tanh");
   populateOpPatterns<math::TanOp>(converter, patterns, "__nv_tanf", "__nv_tan");
 }
 
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
-mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
-  return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth);
+mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth, bool hasRedux) {
+  return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth, hasRedux);
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

-// RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -convert-gpu-to-nvvm='index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s
+// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1' -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -convert-gpu-to-nvvm='has-redux=1 index-bitwidth=32' -split-input-file | FileCheck --check-prefix=CHECK32 %s
 
 gpu.module @test_module {
   // CHECK-LABEL: func @gpu_index_ops()
   // CHECK32-LABEL: func @gpu_index_ops()
   func.func @gpu_index_ops()
       -> (index, index, index, index, index, index,
           index, index, index, index, index, index,
           index) {
gpu.func @test_printf(%arg0: i32, %arg1: f32) {
     // CHECK-NEXT: llvm.store %[[EXT]], %[[EL1]] : !llvm.ptr<f64>
     // CHECK-NEXT: %[[ARGPTR:.*]] = llvm.bitcast %[[ALLOC]] : !llvm.ptr<struct<(i32, f64)>> to !llvm.ptr<i8>
     // CHECK-NEXT: llvm.call @vprintf(%[[FORMATSTART]], %[[ARGPTR]]) : (!llvm.ptr<i8>, !llvm.ptr<i8>) -> i32
     gpu.printf "Hello: %d\n" %arg0, %arg1 : i32, f32
     gpu.return
   }
 }
 
+// -----
+
+gpu.module @test_module {
+  // CHECK-LABEL: func @subgroup_reduce_add
+  gpu.func @subgroup_reduce_add(%arg0 : i32) {
+    // CHECK: nvvm.redux.sync add {{.*}}
+    %result = gpu.subgroup_reduce add %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+  // CHECK-LABEL: func @subgroup_reduce_and
+  gpu.func @subgroup_reduce_and(%arg0 : i32) {
+    // CHECK: nvvm.redux.sync and {{.*}}
+    %result = gpu.subgroup_reduce and %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+  // CHECK-LABEL:  @subgroup_reduce_max
+  gpu.func @subgroup_reduce_max(%arg0 : i32) {
+    // CHECK: nvvm.redux.sync max {{.*}}
+    %result = gpu.subgroup_reduce max %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+  // CHECK-LABEL: @subgroup_reduce_min
+  gpu.func @subgroup_reduce_min(%arg0 : i32) {
+    // CHECK: nvvm.redux.sync min {{.*}}
+    %result = gpu.subgroup_reduce min %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+  // CHECK-LABEL:  @subgroup_reduce_or
+  gpu.func @subgroup_reduce_or(%arg0 : i32) {
+    // CHECK: nvvm.redux.sync or {{.*}}
+    %result = gpu.subgroup_reduce or %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+  // CHECK-LABEL: @subgroup_reduce_xor
+  gpu.func @subgroup_reduce_xor(%arg0 : i32) {
+    // CHECK nvvm.redux.sync xor {{.*}}
+    %result = gpu.subgroup_reduce xor %arg0 uniform {} : (i32) -> (i32)
+    gpu.return
+  }
+}
+