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

[raikonenfnu]

Summary

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

[antiagainst]

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

73

Could you put some comments here like "// We could use the initializer directly; but certain driver compilers have bugs dealing with that. So for now use spv.Store for initialization." so that later we can recall why such behavior?

[raikonenfnu]

Add in code comments to clarify why we need spv.store

[raikonenfnu]

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

73

Could you put some comments here like "// We could use the initializer directly; but certain driver compilers have bugs dealing with that. So for now use spv.Store for initialization." so that later we can recall why such behavior?

Done! thanks :)

[kuhar]

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

77

This variable is unused:

llvm-project/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp:77:22: warning: unused variable 'storeOp' [-Wunused-variable]
      spirv::StoreOp storeOp =
                     ^

[kuhar]

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

77

I fixed this when pushing

[raikonenfnu]

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

77

awesome! thanks 😄

mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp

for (int i = rank - 2; i >= 0; --i) {
       strides[i] = strides[i + 1] * tensorType.getDimSize(i + 1);
     }
 
     Type varType = spirv::PointerType::get(adaptor.getTensor().getType(),
                                            spirv::StorageClass::Function);
 
     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.
       return failure();
     }
 
     auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
     auto indexType = typeConverter.getIndexType();

mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir

 // RUN: mlir-opt -split-input-file -convert-tensor-to-spirv -verify-diagnostics %s | FileCheck %s
 
 //===----------------------------------------------------------------------===//
 // tensor.extract
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func @tensor_extract_constant
 // CHECK-SAME: (%[[A:.+]]: i32, %[[B:.+]]: i32, %[[C:.+]]: i32)
 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
   // CHECK: %[[ADD0:.+]] = spv.IAdd %[[C0]], %[[MUL0]] : i32
   // CHECK: %[[C3:.+]] = spv.Constant 3 : i32
   // CHECK: %[[MUL1:.+]] = spv.IMul %[[C3]], %[[B]] : i32
   // CHECK: %[[ADD1:.+]] = spv.IAdd %[[ADD0]], %[[MUL1]] : i32
   // CHECK: %[[C1:.+]] = spv.Constant 1 : i32
   // CHECK: %[[MUL2:.+]] = spv.IMul %[[C1]], %[[C]] : i32
   // CHECK: %[[ADD2:.+]] = spv.IAdd %[[ADD1]], %[[MUL2]] : i32
   // CHECK: %[[AC:.+]] = spv.AccessChain %[[VAR]][%[[ADD2]]]
   // CHECK: %[[VAL:.+]] = spv.Load "Function" %[[AC]] : i32
   %extract = tensor.extract %cst[%a, %b, %c] : tensor<2x2x3xi32>
   // CHECK: spv.ReturnValue %[[VAL]]
   return %extract : i32
 }