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

[georgemitenkov]

Summary

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.

[georgemitenkov]

Combined conversion patterns for spv.Load and spv.Store into one templated class.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

 /// Creates LLVM dialect constant with the given value.
 static Value createI32ConstantOf(Location loc, PatternRewriter &rewriter,
                                  LLVMTypeConverter &converter, unsigned value) {
   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) {
+  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
 //===----------------------------------------------------------------------===//
 
 /// Converts SPIR-V array type to LLVM array. There is no modelling of array
 /// stride at the moment.
 static Optional<Type> convertArrayType(spirv::ArrayType type,
                                        TypeConverter &converter) {
matchAndRewrite(SPIRVOp operation, ArrayRef<Value> operands,
     rewriter.template replaceOpWithNewOp<LLVM::ICmpOp>(
         operation, dstType,
         rewriter.getI64IntegerAttr(static_cast<int64_t>(predicate)),
         operation.operand1(), operation.operand2());
     return success();
   }
 };
 
+/// 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> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(SPIRVOp notOp, ArrayRef<Value> operands,
patterns.insert<
       // Logical ops
       DirectConversionPattern<spirv::LogicalAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::LogicalOrOp, LLVM::OrOp>,
       IComparePattern<spirv::LogicalEqualOp, LLVM::ICmpPredicate::eq>,
       IComparePattern<spirv::LogicalNotEqualOp, LLVM::ICmpPredicate::ne>,
       NotPattern<spirv::LogicalNotOp>,
 
       // Memory ops
+      LoadStorePattern<spirv::LoadOp>, LoadStorePattern<spirv::StoreOp>,
       VariablePattern,
 
       // Miscellaneous ops
       DirectConversionPattern<spirv::SelectOp, LLVM::SelectOp>,
       DirectConversionPattern<spirv::UndefOp, LLVM::UndefOp>,
 
       // Shift ops
       ShiftPattern<spirv::ShiftRightArithmeticOp, LLVM::AShrOp>,

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

 // 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
 //===----------------------------------------------------------------------===//
 
 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>
   // CHECK: %[[SIZE2:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32