[MLIR][SPIRVToLLVM] Added support of volatile and nontemporal memory access in load/store

[georgemitenkov]

Summary

This patch adds support of Volatile and Nontemporal
memory accesses to spv.Load and spv.Store. These attributes are
modelled with a volatile and nontemporal flags.

[georgemitenkov]

Fixed typo in replaceWithLoadOrStore().

[antiagainst]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

602

Use /*isNonTemporal=*/ to be consistent

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

return rewriter.create<LLVM::ConstantOp>(
       loc, LLVM::LLVMType::getInt32Ty(converter.getDialect()),
       rewriter.getIntegerAttr(rewriter.getI32Type(), value));
 }
 
 /// Utility for `spv.Load` and `spv.Store` conversion.
 static LogicalResult replaceWithLoadOrStore(Operation *op,
                                             ConversionPatternRewriter &rewriter,
                                             LLVMTypeConverter &typeConverter,
-                                            unsigned alignment) {
+                                            unsigned alignment, bool isVolatile,
+                                            bool isNonTemporal) {
   if (auto loadOp = dyn_cast<spirv::LoadOp>(op)) {
     auto dstType = typeConverter.convertType(loadOp.getType());
     if (!dstType)
       return failure();
-    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(loadOp, dstType, loadOp.ptr(),
-                                              alignment);
+    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(
+        loadOp, dstType, loadOp.ptr(), alignment, isVolatile, isNonTemporal);
     return success();
   }
   auto storeOp = cast<spirv::StoreOp>(op);
   rewriter.replaceOpWithNewOp<LLVM::StoreOp>(storeOp, storeOp.value(),
-                                             storeOp.ptr(), alignment);
+                                             storeOp.ptr(), alignment,
+                                             isVolatile, isNonTemporal);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//
 
 /// Converts SPIR-V array type to LLVM array. There is no modelling of array

 template <typename SPIRVop>
 class LoadStorePattern : public SPIRVToLLVMConversion<SPIRVop> {
 public:
   using SPIRVToLLVMConversion<SPIRVop>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(SPIRVop op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    if (op.memory_access().hasValue() &&
-        op.memory_access().getValue() != spirv::MemoryAccess::None) {
-      auto memoryAccess = op.memory_access().getValue();
-      if (memoryAccess == spirv::MemoryAccess::Aligned) {
-        unsigned alignment = op.alignment().getValue().getZExtValue();
-        replaceWithLoadOrStore(op, rewriter, this->typeConverter, alignment);
-        return success();
-      }
+
+    if (!op.memory_access().hasValue()) {
+      replaceWithLoadOrStore(op, rewriter, this->typeConverter, /*alignment=*/0,
+                             /*isVolatile=*/false, /*isNonTemporal=*/ false);
+      return success();
+    }
+    switch (memoryAccess) {
+    case spirv::MemoryAccess::Aligned:
+    case spirv::MemoryAccess::None:
+    case spirv::MemoryAccess::Nontemporal:
+    case spirv::MemoryAccess::Volatile: {
+      unsigned alignment = memoryAccess == spirv::MemoryAccess::Aligned
+                               ? op.alignment().getValue().getZExtValue()
+                               : 0;
+      bool isNonTemporal = memoryAccess == spirv::MemoryAccess::Nontemporal;
+      bool isVolatile = memoryAccess == spirv::MemoryAccess::Volatile;
+      replaceWithLoadOrStore(op, rewriter, this->typeConverter, alignment,
+                             isVolatile, isNonTemporal);
+    default:
       // There is no support of other memory access attributes.
       return failure();
     }
-    replaceWithLoadOrStore(op, rewriter, this->typeConverter, 0);
-    return success();
   }
 };
 
 /// Converts `spv.Not` and `spv.LogicalNot` into LLVM dialect.
 template <typename SPIRVOp>
 class NotPattern : public SPIRVToLLVMConversion<SPIRVOp> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;

mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

 
 func @load_with_alignment() {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   // CHECK: %{{.*}} = llvm.load %{{.*}} {alignment = 4 : i64} : !llvm<"float*">
   %1 = spv.Load "Function" %0 ["Aligned", 4] : f32
   return
 }
 
+func @load_volatile() {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: %{{.*}} = llvm.load volatile %{{.*}} : !llvm<"float*">
+  %1 = spv.Load "Function" %0 ["Volatile"] : f32
+  return
+}
+
+func @load_nontemporal() {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: %{{.*}} = llvm.load %{{.*}} {nontemporal} : !llvm<"float*">
+  %1 = spv.Load "Function" %0 ["Nontemporal"] : f32
+  return
+}
+
 //===----------------------------------------------------------------------===//
 // spv.Store
 //===----------------------------------------------------------------------===//
 
 func @store(%arg0 : f32) -> () {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   // CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm<"float*">
   spv.Store "Function" %0, %arg0 : f32

 
 func @store_with_alignment(%arg0 : f32) -> () {
   %0 = spv.Variable : !spv.ptr<f32, Function>
   // CHECK: llvm.store %{{.*}}, %{{.*}} {alignment = 4 : i64} : !llvm<"float*">
   spv.Store "Function" %0, %arg0 ["Aligned", 4] : f32
   return
 }
 
+func @store_volatile(%arg0 : f32) -> () {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: llvm.store volatile %{{.*}}, %{{.*}} : !llvm<"float*">
+  spv.Store "Function" %0, %arg0 ["Volatile"] : f32
+  return
+}
+
+func @store_nontemporal(%arg0 : f32) -> () {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: llvm.store %{{.*}}, %{{.*}} {nontemporal} : !llvm<"float*">
+  spv.Store "Function" %0, %arg0 ["Nontemporal"] : f32
+  return
+}
+
 //===----------------------------------------------------------------------===//
 // spv.Variable
 //===----------------------------------------------------------------------===//
 
 func @variable_scalar() {
   // CHECK: %[[SIZE1:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
   // CHECK: %{{.*}} = llvm.alloca %[[SIZE1]] x !llvm.float : (!llvm.i32) -> !llvm<"float*">
   %0 = spv.Variable : !spv.ptr<f32, Function>