CSE spirv Input loads (#172)

Author: Hardcode84

mlir/lib/dialect/plier_util/dialect.cpp

@@ -19,6 +19,7 @@
 #include <mlir/IR/BuiltinOps.h>
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/DialectImplementation.h>
+#include <mlir/IR/Dominance.h>
 #include <mlir/IR/PatternMatch.h>
 #include <mlir/Transforms/InliningUtils.h>
 
@@ -28,6 +29,7 @@
 #include <mlir/Dialect/Linalg/IR/Linalg.h>
 #include <mlir/Dialect/MemRef/IR/MemRef.h>
 #include <mlir/Dialect/SCF/SCF.h>
+#include <mlir/Dialect/SPIRV/IR/SPIRVOps.h>
 
 #include <llvm/ADT/SmallBitVector.h>
 #include <llvm/ADT/TypeSwitch.h>
@@ -185,13 +187,52 @@ struct FillExtractSlice
     return mlir::success();
   }
 };
+
+struct SpirvInputCSE : public mlir::OpRewritePattern<mlir::spirv::LoadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(mlir::spirv::LoadOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    auto ptr = op.ptr();
+    if (ptr.getType().cast<mlir::spirv::PointerType>().getStorageClass() !=
+        mlir::spirv::StorageClass::Input)
+      return mlir::failure();
+
+    auto func = op->getParentOfType<mlir::spirv::FuncOp>();
+    if (!func)
+      return mlir::failure();
+
+    mlir::DominanceInfo dom;
+    mlir::spirv::LoadOp prevLoad;
+    func->walk([&](mlir::spirv::LoadOp load) -> mlir::WalkResult {
+      if (load == op)
+        return mlir::WalkResult::interrupt();
+
+      if (load->getOperands() == op->getOperands() &&
+          load->getResultTypes() == op->getResultTypes() &&
+          dom.properlyDominates(load.getOperation(), op)) {
+        prevLoad = load;
+        return mlir::WalkResult::interrupt();
+      }
+
+      return mlir::WalkResult::advance();
+    });
+
+    if (!prevLoad)
+      return mlir::failure();
+
+    rewriter.replaceOp(op, prevLoad.getResult());
+    return mlir::success();
+  }
+};
 } // namespace
 
 void PlierUtilDialect::getCanonicalizationPatterns(
     mlir::RewritePatternSet &results) const {
   results.add<DimExpandShape<mlir::tensor::DimOp, mlir::tensor::ExpandShapeOp>,
               DimExpandShape<mlir::memref::DimOp, mlir::memref::ExpandShapeOp>,
-              DimInsertSlice, FillExtractSlice>(getContext());
+              DimInsertSlice, FillExtractSlice, SpirvInputCSE>(getContext());
 }
 
 OpaqueType OpaqueType::get(mlir::MLIRContext *context) {

numba_dpcomp/numba_dpcomp/mlir/tests/test_gpu.py

@@ -654,6 +654,36 @@ def func(a, b, c, res):
         assert_equal(gpu_res, sim_res)
 
 
+@require_gpu
+def test_input_load_cse():
+    def func(c):
+        i = get_global_id(0)
+        j = get_global_id(1)
+        k = get_global_id(2)
+        c[i, j, k] = i + 10 * j + 100 * k
+
+    sim_func = kernel_sim(func)
+    gpu_func = kernel_cached(func)
+
+    a = np.array([[[1, 2, 3], [4, 5, 6]]], np.float32)
+    sim_res = np.zeros(a.shape, a.dtype)
+    sim_func[a.shape, DEFAULT_LOCAL_SIZE](sim_res)
+
+    gpu_res = np.zeros(a.shape, a.dtype)
+
+    with print_pass_ir(["SerializeSPIRVPass"], []):
+        gpu_func[a.shape, DEFAULT_LOCAL_SIZE](gpu_res)
+        ir = get_print_buffer()
+        assert (
+            ir.count(
+                'spv.Load "Input" %__builtin_var_GlobalInvocationId___addr : vector<3xi64>'
+            )
+            == 1
+        ), ir
+
+    assert_equal(gpu_res, sim_res)
+
+
 @require_dpctl
 def test_dpctl_simple1():
     def func(a, b, c):