[Mlir-commits] [mlir] f148f3b - [mlir][spirv] Use spv.store instead of init in var

[Jakub Kuderski]

Author: Stanley Winata
Date: 2022-08-22T23:41:27-04:00
New Revision: f148f3b2e915eece0a6c0bad13d31ab25e635d8d

URL: https://github.com/llvm/llvm-project/commit/f148f3b2e915eece0a6c0bad13d31ab25e635d8d
DIFF: https://github.com/llvm/llvm-project/commit/f148f3b2e915eece0a6c0bad13d31ab25e635d8d.diff

LOG: [mlir][spirv] Use spv.store instead of init in var

- Add spv.store instead of init for spv.variable to fix data issues in
  some GPU drivers

Reviewed By: antiagainst
Patch By: raikonenfnu

Differential Revision: https://reviews.llvm.org/D132427

Added: 
    

Modified: 
    mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
    mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
index 44ba0d0adaab4..9421c3ecd7625 100644
--- a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
+++ b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
@@ -68,9 +68,13 @@ class TensorExtractPattern final
 
     spirv::VariableOp varOp;
     if (adaptor.getTensor().getDefiningOp<spirv::ConstantOp>()) {
-      varOp = rewriter.create<spirv::VariableOp>(
-          loc, varType, spirv::StorageClass::Function,
-          /*initializer=*/adaptor.getTensor());
+      // We could use the initializer directly; but certain driver compilers
+      // have bugs dealing with that. So for now, use spv.Store for
+      // initialization.
+      varOp = rewriter.create<spirv::VariableOp>(loc, varType,
+                                                 spirv::StorageClass::Function,
+                                                 /*initializer=*/nullptr);
+      rewriter.create<spirv::StoreOp>(loc, varOp, adaptor.getTensor());
     } else {
       // Need to store the value to the local variable. It's questionable
       // whether we want to support such case though.

diff  --git a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
index cbec0f5af9946..1ad8df1f93e44 100644
--- a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
@@ -9,7 +9,8 @@
 func.func @tensor_extract_constant(%a : index, %b: index, %c: index) -> i32 {
   // CHECK: %[[CST:.+]] = spv.Constant dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]>
   %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]> : tensor<2x2x3xi32>
-  // CHECK: %[[VAR:.+]] = spv.Variable init(%[[CST]]) : !spv.ptr<!spv.array<12 x i32>, Function>
+  // CHECK: %[[VAR:.+]] = spv.Variable : !spv.ptr<!spv.array<12 x i32>, Function>
+  // CHECK: spv.Store "Function" %[[VAR]], %[[CST]] : !spv.array<12 x i32>
   // CHECK: %[[C0:.+]] = spv.Constant 0 : i32
   // CHECK: %[[C6:.+]] = spv.Constant 6 : i32
   // CHECK: %[[MUL0:.+]] = spv.IMul %[[C6]], %[[A]] : i32