[mlir][AMDGPU] Improving DPP implementation of subgrroup reduce

Signed-off-by: Muzammiluddin Syed <[email protected]>

Author: Muzammiluddin-Syed-ECE

mlir/include/mlir/Dialect/GPU/Transforms/Passes.h

@@ -74,13 +74,15 @@ void populateGpuLowerClusteredSubgroupReduceToShufflePatterns(
 /// `subgroupSize` lanes. Applicable only to AMD GPUs.
 void populateGpuLowerSubgroupReduceToDPPPatterns(RewritePatternSet &patterns,
                                                  unsigned subgroupSize,
+                                                 unsigned shuffleBitwidth,
                                                  amdgpu::Chipset chipset,
                                                  PatternBenefit benefit = 1);
 
 /// Disjoint counterpart of `populateGpuLowerSubgroupReduceToDPPPatterns`
 /// that only matches `gpu.subgroup_reduce` ops with a `cluster_size`.
 void populateGpuLowerClusteredSubgroupReduceToDPPPatterns(
-    RewritePatternSet &patterns, unsigned subgroupSize, amdgpu::Chipset chipset,
+    RewritePatternSet &patterns, unsigned subgroupSize,
+    unsigned shuffleBitwidth, amdgpu::Chipset chipset,
     PatternBenefit benefit = 1);
 
 /// Collect all patterns to rewrite ops within the GPU dialect.

mlir/include/mlir/Dialect/GPU/Utils/GPUUtils.h

@@ -29,6 +29,8 @@ class LaunchOp;
 
 /// Returns the matching vector combining kind.
 vector::CombiningKind convertReductionKind(gpu::AllReduceOperation mode);
+/// Returns the matching gpu allreduce mode.
+gpu::AllReduceOperation convertReductionMode(vector::CombiningKind kind);
 } // namespace gpu
 
 /// Get a gpu.func created from outlining the region of a gpu.launch op with the

mlir/include/mlir/Dialect/GPU/Utils/ReductionUtils.h

@@ -0,0 +1,41 @@
+//===- ReductionUtils.h - Reduction Utilities -------------*- 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_DIALECT_GPU_TRANSFORMS_REDUCTIONUTILS_H_
+#define MLIR_DIALECT_GPU_TRANSFORMS_REDUCTIONUTILS_H_
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
+#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+
+struct ClusterInfo {
+  unsigned clusterStride;
+  unsigned clusterSize;
+  unsigned subgroupSize;
+};
+
+FailureOr<ClusterInfo> getAndValidateClusterInfo(gpu::SubgroupReduceOp op,
+  unsigned subgroupSize);
+
+FailureOr<Value>
+createSubgroupDPPReduction(PatternRewriter &rewriter, gpu::SubgroupReduceOp op,
+                           Value input, gpu::AllReduceOperation mode,
+                           const ClusterInfo &ci, amdgpu::Chipset chipset,
+                           function_ref<Value(Value)> packFn,
+                           function_ref<Value(Value)> unpackFn);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_GPU_TRANSFORMS_REDUCTIONUTILS_H_

mlir/lib/Dialect/GPU/CMakeLists.txt

@@ -53,6 +53,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
 
   LINK_LIBS PUBLIC
   MLIRAffineUtils
+  MLIRAMDGPUDialect
   MLIRArithDialect
   MLIRAsyncDialect
   MLIRBufferizationDialect