changing from dialect converstion to gpu transform pass

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

Author: Muzammiluddin-Syed-ECE

mlir/include/mlir/Conversion/Passes.h

@@ -34,7 +34,6 @@
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRVPass.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
-#include "mlir/Conversion/GPUToAMDGPU/GPUToAMDGPU.h"
 #include "mlir/Conversion/GPUToLLVMSPV/GPUToLLVMSPVPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"

mlir/include/mlir/Conversion/Passes.td

@@ -643,22 +643,6 @@ def ConvertGPUToSPIRV : Pass<"convert-gpu-to-spirv", "ModuleOp"> {
   ];
 }
 
-//===----------------------------------------------------------------------===//
-// GPUToAMDGPU
-//===----------------------------------------------------------------------===//
-
-def ConvertGPUToAMDGPUPass : Pass<"convert-gpu-to-amdgpu"> {
-  let summary = "Generate AMDGPU operations for gpu operations";
-  let dependentDialects = [
-    "LLVM::LLVMDialect",
-    "::mlir::gpu::GPUDialect",
-    "amdgpu::AMDGPUDialect",
-  ];
-  let options = [Option<"subgroupSize", "subgroup-size", "unsigned",
-                        /*default=*/"64",
-                        "Size of subgroup">];
-}
-
 //===----------------------------------------------------------------------===//
 // ConvertIndexToLLVMPass
 //===----------------------------------------------------------------------===//

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

@@ -94,6 +94,9 @@ void populateGpuDecomposeMemrefsPatterns(RewritePatternSet &patterns);
 /// Erase barriers that do not enforce conflicting memory side effects.
 void populateGpuEliminateBarriersPatterns(RewritePatternSet &patterns);
 
+void populateGpuDecomposeSubgroupReduceToDppPatterns(
+    RewritePatternSet &patterns, unsigned subgroupSize);
+
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/GPU/Transforms/Passes.h.inc"

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

@@ -258,4 +258,28 @@ def GpuSPIRVAttachTarget: Pass<"spirv-attach-target", ""> {
   ];
 }
 
+def GpuDecomposeSubgroupReduceToDppPass: Pass<"decompose-subgroup-reduce-to-dpp", ""> {
+  let summary = "Decomposes reductions across subgroups to data parallel primitives.";
+  let description = [{
+    This pass decomposes gpu::subgroup_reduce ops to AMDGPU::DPP ops where
+    applicable, while respecting cluster size and stride. 
+
+    Example:
+    ```
+    // File: in.mlir:
+    [TO-DO]
+    // mlir-opt --decompose-subgroup-reduce-to-dpp subgroup-size=64 in.mlir
+    [TO-DO]
+    ```
+  }];
+  let dependentDialects = [
+    "amdgpu::AMDGPUDialect",
+    "LLVM::LLVMDialect",
+    "ROCDL::ROCDLDialect",
+  ];
+  let options = [Option<"subgroupSize", "subgroup-size", "unsigned",
+                        /*default=*/"64",
+                        "Size of subgroup">];
+}
+
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -1333,7 +1333,6 @@ struct ConvertAMDGPUToROCDLPass
   using Base::Base;
 
   void runOnOperation() override {
-    llvm::errs() << " WHEN DOES AMDGPU TO ROCDL RUN\n";
     MLIRContext *ctx = &getContext();
     FailureOr<Chipset> maybeChipset = Chipset::parse(chipset);
     if (failed(maybeChipset)) {

mlir/lib/Conversion/CMakeLists.txt

@@ -24,7 +24,6 @@ add_subdirectory(FuncToEmitC)
 add_subdirectory(FuncToLLVM)
 add_subdirectory(FuncToSPIRV)
 add_subdirectory(GPUCommon)
-add_subdirectory(GPUToAMDGPU)
 add_subdirectory(GPUToLLVMSPV)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)

mlir/lib/Conversion/GPUToAMDGPU/CMakeLists.txt

@@ -15,7 +15,6 @@ add_mlir_conversion_library(MLIRGPUToROCDLTransforms
   MLIRMathToLLVM
   MLIRMathToROCDL
   MLIRAMDGPUToROCDL
-  MLIRGPUToAMDGPU
   MLIRFuncToLLVM
   MLIRGPUDialect
   MLIRGPUToGPURuntimeTransforms

mlir/lib/Conversion/GPUToROCDL/CMakeLists.txt

@@ -41,6 +41,7 @@ add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/ShuffleRewriter.cpp
   Transforms/SPIRVAttachTarget.cpp
   Transforms/SubgroupReduceLowering.cpp
+  Transforms/DecomposeSubgroupReduceToDpp.cpp
 
   OBJECT
 

mlir/lib/Dialect/GPU/CMakeLists.txt

@@ -1,16 +1,17 @@
-//===- GPUToAMDGPU.cpp - GPU to AMDGPU dialect conversion -------===//
+//===- DecomposeSubgroupReduceToDPP.cpp - Decompose subgroup reduce pass -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
+//
+// This file implements decompose subgroup reduce to DPP pass.
+//
+//===----------------------------------------------------------------------===//
 
-#include "mlir/Conversion/GPUToAMDGPU/GPUToAMDGPU.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
 
-#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
-#include "mlir/Conversion/LLVMCommon/Pattern.h"
-#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
 #include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -22,20 +23,14 @@
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Dialect/GPU/Transforms/Passes.h"
 
 #include "llvm/Support/FormatVariadic.h"
-#include "llvm/Support/MathExtras.h"
-#include <cassert>
-#include <cstdint>
-
-#include "../LLVMCommon/MemRefDescriptor.h"
-
-#include "llvm/ADT/STLExtras.h"
-#include <optional>
 
 namespace mlir {
-#define GEN_PASS_DEF_CONVERTGPUTOAMDGPUPASS
-#include "mlir/Conversion/Passes.h.inc"
+#define GEN_PASS_DEF_GPUDECOMPOSESUBGROUPREDUCETODPPPASS
+#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
@@ -144,8 +139,8 @@ Value createSubgroupDPPReduction(OpBuilder &b, Location loc, Value input,
 struct ScalarSubgroupReduceToShuffles final
     : OpRewritePattern<gpu::SubgroupReduceOp> {
   ScalarSubgroupReduceToShuffles(MLIRContext *ctx, unsigned subgroupSize,
-                                 bool matchClustered, PatternBenefit benefit)
-      : OpRewritePattern(ctx, benefit), subgroupSize(subgroupSize),
+                                 bool matchClustered)
+      : OpRewritePattern(ctx), subgroupSize(subgroupSize),
         matchClustered(matchClustered) {}
 
   LogicalResult matchAndRewrite(gpu::SubgroupReduceOp op,
@@ -174,30 +169,24 @@ struct ScalarSubgroupReduceToShuffles final
   bool matchClustered = false;
 };
 
-struct ConvertGPUToAMDGPUPass
-    : public impl::ConvertGPUToAMDGPUPassBase<ConvertGPUToAMDGPUPass> {
+struct GpuDecomposeSubgroupReduceToDppPass
+    : public impl::GpuDecomposeSubgroupReduceToDppPassBase<
+          GpuDecomposeSubgroupReduceToDppPass> {
   using Base::Base;
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
-    LLVMTypeConverter converter(&getContext());
-    LLVMConversionTarget target(getContext());
-    target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
-    target.addLegalDialect<::mlir::amdgpu::AMDGPUDialect>();
-    target.addLegalDialect<::mlir::ROCDL::ROCDLDialect>();
-
-    int subgroupSizeInt = static_cast<int>(subgroupSize);
-    populateSubgroupReduceLoweringPatterns(converter, patterns, subgroupSizeInt,
-                                           PatternBenefit(1));
-    if (failed(applyPartialConversion(getOperation(), target,
-                                      std::move(patterns))))
-      signalPassFailure();
+    // int subgroupSizeInt = static_cast<int>(subgroupSize);
+    populateGpuDecomposeSubgroupReduceToDppPatterns(patterns, subgroupSize);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+      return signalPassFailure();
   }
 };
+
 } // namespace
 
-void mlir::populateSubgroupReduceLoweringPatterns(
-    LLVMTypeConverter &converter, RewritePatternSet &patterns, unsigned subgroupSize, PatternBenefit benefit) {
+void mlir::populateGpuDecomposeSubgroupReduceToDppPatterns(
+    RewritePatternSet &patterns, unsigned subgroupSize) {
   patterns.add<ScalarSubgroupReduceToShuffles>(
-      patterns.getContext(), subgroupSize, /*matchClustered=*/true, benefit);
-}
+      patterns.getContext(), subgroupSize, /*matchClustered=*/true);
+}