[Mlir-commits] [mlir] 8be0371 - [MLIR][SPIRVToLLVM] Conversion of load and store SPIR-V ops

[George Mitenkov]

Author: George Mitenkov
Date: 2020-07-24T16:31:45+03:00
New Revision: 8be0371eb79fb25ba22846410e1a31c88b4c89c2

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

LOG: [MLIR][SPIRVToLLVM] Conversion of load and store SPIR-V ops

This patch introduces conversion pattern for `spv.Store` and `spv.Load`.
Only op with `Function` Storage Class is supported at the moment
because `spv.GlobalVariable` has not been introduced yet. If the op
has memory access attribute, then there are the following cases.
If the access is `Aligned`, add alignment to the op builder. Otherwise
the conversion fails as other cases are not supported yet because they
require additional attributes for `llvm.store`/`llvm.load` ops: e.g.
`volatile` and `nontemporal`.

Reviewed By: antiagainst

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

Added: 
    

Modified: 
    mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
    mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp b/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
index cb4898c346a6..12aa6573ef37 100644
--- a/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
+++ b/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
@@ -185,6 +185,25 @@ static Value createI32ConstantOf(Location loc, PatternRewriter &rewriter,
       rewriter.getIntegerAttr(rewriter.getI32Type(), value));
 }
 
+/// Utility for `spv.Load` and `spv.Store` conversion.
+static LogicalResult replaceWithLoadOrStore(Operation *op,
+                                            ConversionPatternRewriter &rewriter,
+                                            LLVMTypeConverter &typeConverter,
+                                            unsigned alignment) {
+  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);
+    return success();
+  }
+  auto storeOp = cast<spirv::StoreOp>(op);
+  rewriter.replaceOpWithNewOp<LLVM::StoreOp>(storeOp, storeOp.value(),
+                                             storeOp.ptr(), alignment);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//
@@ -566,6 +585,31 @@ class IComparePattern : public SPIRVToLLVMConversion<SPIRVOp> {
   }
 };
 
+/// Converts `spv.Load` and `spv.Store` to LLVM dialect.
+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();
+      }
+      // 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> {
@@ -973,6 +1017,7 @@ void mlir::populateSPIRVToLLVMConversionPatterns(
       NotPattern<spirv::LogicalNotOp>,
 
       // Memory ops
+      LoadStorePattern<spirv::LoadOp>, LoadStorePattern<spirv::StoreOp>,
       VariablePattern,
 
       // Miscellaneous ops

diff  --git a/mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir b/mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir
index 4c549ab8d619..362c702e8d3c 100644
--- a/mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir
+++ b/mlir/test/Conversion/SPIRVToLLVM/memory-ops-to-llvm.mlir
@@ -1,5 +1,55 @@
 // RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
 
+//===----------------------------------------------------------------------===//
+// spv.Load
+//===----------------------------------------------------------------------===//
+
+func @load() {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  //  CHECK: %{{.*}} = llvm.load %{{.*}} : !llvm<"float*">
+  %1 = spv.Load "Function" %0 : f32
+  return
+}
+
+func @load_none() {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  //  CHECK: %{{.*}} = llvm.load %{{.*}} : !llvm<"float*">
+  %1 = spv.Load "Function" %0 ["None"] : f32
+  return
+}
+
+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
+}
+
+//===----------------------------------------------------------------------===//
+// spv.Store
+//===----------------------------------------------------------------------===//
+
+func @store(%arg0 : f32) -> () {
+  %0 = spv.Variable : !spv.ptr<f32, Function>
+  // CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm<"float*">
+  spv.Store "Function" %0, %arg0 : f32
+  return
+}
+
+func @store_composite(%arg0 : !spv.struct<f64>) -> () {
+  %0 = spv.Variable : !spv.ptr<!spv.struct<f64>, Function>
+  // CHECK: llvm.store %{{.*}}, %{{.*}} : !llvm<"<{ double }>*">
+  spv.Store "Function" %0, %arg0 : !spv.struct<f64>
+  return
+}
+
+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
+}
+
 //===----------------------------------------------------------------------===//
 // spv.Variable
 //===----------------------------------------------------------------------===//