[python][gpu] GPU code refactoring (#207)

Author: Hardcode84

mlir/lib/Conversion/gpu_runtime_to_llvm.cpp

@@ -498,8 +498,10 @@ class ConvertGpuKernelLaunchPattern
         mlir::MemRefDescriptor desc(kernelParams[i]);
         if (memrefType.getMemorySpace() == localMemStorageClass) {
           auto rank = static_cast<unsigned>(memrefType.getRank());
+          auto typeSize = std::max(memrefType.getElementTypeBitWidth(), 8u) / 8;
           mlir::Value size = rewriter.create<mlir::LLVM::ConstantOp>(
-              loc, llvmIndexType, rewriter.getIntegerAttr(llvmIndexType, 0));
+              loc, llvmIndexType,
+              rewriter.getIntegerAttr(llvmIndexType, typeSize));
           for (auto i : llvm::seq(0u, rank)) {
             auto dim = desc.size(rewriter, loc, i);
             size = rewriter.create<mlir::LLVM::MulOp>(loc, llvmIndexType, size,

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

@@ -1167,6 +1167,9 @@ class ConvertBarrierOps : public mlir::OpRewritePattern<mlir::func::CallOp> {
   mlir::LogicalResult
   matchAndRewrite(mlir::func::CallOp op,
                   mlir::PatternRewriter &rewriter) const override {
+    if (!op->getParentOfType<mlir::gpu::LaunchOp>())
+      return mlir::failure();
+
     auto operands = op.operands();
     if (operands.size() != 1)
       return mlir::failure();
@@ -1197,6 +1200,10 @@ class ConvertBarrierOps : public mlir::OpRewritePattern<mlir::func::CallOp> {
   }
 };
 
+struct LowerGpuBuiltins2Pass
+    : public plier::RewriteWrapperPass<LowerGpuBuiltins2Pass, void, void,
+                                       ConvertBarrierOps> {};
+
 class ConvertArrayAllocOps : public mlir::OpRewritePattern<mlir::func::CallOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
@@ -1279,9 +1286,8 @@ class ConvertArrayAllocOps : public mlir::OpRewritePattern<mlir::func::CallOp> {
   }
 };
 
-struct LowerGpuBuiltins2Pass
-    : public plier::RewriteWrapperPass<LowerGpuBuiltins2Pass, void, void,
-                                       ConvertBarrierOps,
+struct LowerGpuBuiltins3Pass
+    : public plier::RewriteWrapperPass<LowerGpuBuiltins3Pass, void, void,
                                        ConvertArrayAllocOps> {};
 
 class GpuLaunchSinkOpsPass
@@ -1412,6 +1418,7 @@ static void populateLowerToGPUPipelineLow(mlir::OpPassManager &pm) {
 
   commonOptPasses(funcPM);
   funcPM.addPass(std::make_unique<KernelMemrefOpsMovementPass>());
+  funcPM.addPass(std::make_unique<LowerGpuBuiltins2Pass>());
   funcPM.addPass(std::make_unique<SinkGpuDimsPass>());
   funcPM.addPass(std::make_unique<GpuLaunchSinkOpsPass>());
   pm.addPass(mlir::createGpuKernelOutliningPass());
@@ -1425,7 +1432,7 @@ static void populateLowerToGPUPipelineLow(mlir::OpPassManager &pm) {
       pm.nest<mlir::gpu::GPUModuleOp>().nest<mlir::gpu::GPUFuncOp>();
   gpuFuncPM.addPass(mlir::arith::createArithmeticExpandOpsPass());
   gpuFuncPM.addPass(std::make_unique<FlattenScfPass>());
-  gpuFuncPM.addPass(std::make_unique<LowerGpuBuiltins2Pass>());
+  gpuFuncPM.addPass(std::make_unique<LowerGpuBuiltins3Pass>());
   commonOptPasses(gpuFuncPM);
   gpuFuncPM.addPass(std::make_unique<AssumeGpuIdRangePass>());