GPU lowering refactoring (#157)

Author: Hardcode84

mlir/include/mlir-extensions/transforms/common_opts.hpp

@@ -14,9 +14,12 @@
 
 #pragma once
 
+#include <memory>
+
 namespace mlir {
 class RewritePatternSet;
 class MLIRContext;
+class Pass;
 } // namespace mlir
 
 namespace plier {
@@ -25,4 +28,6 @@ void populateCanonicalizationPatterns(mlir::MLIRContext &context,
 
 void populateCommonOptsPatterns(mlir::MLIRContext &context,
                                 mlir::RewritePatternSet &patterns);
+
+std::unique_ptr<mlir::Pass> createCommonOptsPass();
 } // namespace plier

mlir/include/mlir-extensions/transforms/index_type_propagation.hpp

@@ -20,6 +20,6 @@ class MLIRContext;
 } // namespace mlir
 
 namespace plier {
-void populate_index_propagate_patterns(mlir::MLIRContext &context,
-                                       mlir::RewritePatternSet &patterns);
+void populateIndexPropagatePatterns(mlir::MLIRContext &context,
+                                    mlir::RewritePatternSet &patterns);
 }

mlir/lib/transforms/common_opts.cpp

@@ -19,6 +19,8 @@
 #include "mlir-extensions/transforms/index_type_propagation.hpp"
 #include "mlir-extensions/transforms/loop_rewrites.hpp"
 #include "mlir-extensions/transforms/memory_rewrites.hpp"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include <mlir/Dialect/Math/IR/Math.h>
 #include <mlir/Dialect/MemRef/IR/MemRef.h>
@@ -129,6 +131,20 @@ struct PowSimplify : public mlir::OpRewritePattern<mlir::math::PowFOp> {
     return mlir::failure();
   }
 };
+
+struct CommonOptsPass
+    : public mlir::PassWrapper<CommonOptsPass, mlir::OperationPass<void>> {
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    mlir::RewritePatternSet patterns(ctx);
+
+    plier::populateCommonOptsPatterns(*ctx, patterns);
+
+    (void)mlir::applyPatternsAndFoldGreedily(getOperation(),
+                                             std::move(patterns));
+  }
+};
 } // namespace
 
 void plier::populateCanonicalizationPatterns(
@@ -155,5 +171,9 @@ void plier::populateCommonOptsPatterns(mlir::MLIRContext &context,
       // clang-format on
       >(&context);
 
-  plier::populate_index_propagate_patterns(context, patterns);
+  plier::populateIndexPropagatePatterns(context, patterns);
+}
+
+std::unique_ptr<mlir::Pass> plier::createCommonOptsPass() {
+  return std::make_unique<CommonOptsPass>();
 }

mlir/lib/transforms/index_type_propagation.cpp

@@ -147,8 +147,8 @@ struct CmpIndexCastSimplify
 };
 } // namespace
 
-void plier::populate_index_propagate_patterns(
-    mlir::MLIRContext &context, mlir::RewritePatternSet &patterns) {
+void plier::populateIndexPropagatePatterns(mlir::MLIRContext &context,
+                                           mlir::RewritePatternSet &patterns) {
   patterns
       .insert<CmpIndexCastSimplify, ArithIndexCastSimplify<mlir::arith::SubIOp>,
               ArithIndexCastSimplify<mlir::arith::AddIOp>,

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/lower_to_gpu.cpp

@@ -67,6 +67,7 @@
 #include "mlir-extensions/dialect/plier_util/dialect.hpp"
 #include "mlir-extensions/transforms/call_lowering.hpp"
 #include "mlir-extensions/transforms/cast_utils.hpp"
+#include "mlir-extensions/transforms/common_opts.hpp"
 #include "mlir-extensions/transforms/const_utils.hpp"
 #include "mlir-extensions/transforms/func_utils.hpp"
 #include "mlir-extensions/transforms/pipeline_utils.hpp"
@@ -2986,9 +2987,9 @@ class GpuLaunchSinkOpsPass
 };
 
 static void commonOptPasses(mlir::OpPassManager &pm) {
-  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(plier::createCommonOptsPass());
   pm.addPass(mlir::createCSEPass());
-  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(plier::createCommonOptsPass());
 }
 
 static void populateLowerToGPUPipelineHigh(mlir::OpPassManager &pm) {
@@ -3014,8 +3015,6 @@ static void populateLowerToGPUPipelineLow(mlir::OpPassManager &pm) {
   funcPM.addPass(std::make_unique<UnstrideMemrefsPass>());
   funcPM.addPass(mlir::createLowerAffinePass());
 
-  // TODO: mlir::gpu::GPUModuleOp pass
-  pm.addNestedPass<mlir::FuncOp>(mlir::arith::createArithmeticExpandOpsPass());
   commonOptPasses(funcPM);
   funcPM.addPass(std::make_unique<KernelMemrefOpsMovementPass>());
   funcPM.addPass(std::make_unique<GpuLaunchSinkOpsPass>());

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_linalg.cpp

@@ -642,6 +642,7 @@ computeIndices(mlir::OpBuilder &builder, mlir::Location loc, mlir::Value value,
     auto valType = indexVal.getType();
 
     auto len = getDim(dim);
+    bool ignoreNegativeInd = false;
     auto handleNegativeVal = [&](mlir::OpFoldResult val) -> mlir::Value {
       mlir::Value idx;
       if (auto v = val.dyn_cast<mlir::Value>()) {
@@ -651,11 +652,16 @@ computeIndices(mlir::OpBuilder &builder, mlir::Location loc, mlir::Value value,
         auto attrVal = attr.cast<mlir::IntegerAttr>().getValue().getSExtValue();
         idx = builder.create<mlir::arith::ConstantIndexOp>(loc, attrVal);
       }
-      auto isNeg = builder.createOrFold<mlir::arith::CmpIOp>(
-          loc, mlir::arith::CmpIPredicate::slt, idx, zero);
-      auto negIndex = builder.createOrFold<mlir::arith::AddIOp>(loc, len, idx);
-      return builder.createOrFold<mlir::arith::SelectOp>(loc, isNeg, negIndex,
-                                                         idx);
+      if (ignoreNegativeInd) {
+        return idx;
+      } else {
+        auto isNeg = builder.createOrFold<mlir::arith::CmpIOp>(
+            loc, mlir::arith::CmpIPredicate::slt, idx, zero);
+        auto negIndex =
+            builder.createOrFold<mlir::arith::AddIOp>(loc, len, idx);
+        return builder.createOrFold<mlir::arith::SelectOp>(loc, isNeg, negIndex,
+                                                           idx);
+      }
     };
 
     if (auto sliceType = valType.dyn_cast<plier::SliceType>()) {