[PATCH] D74874: [mlir][spirv] Add lowering for load/store zero-rank memref from std to SPIR-V.

[Han-Chung Wang via Phabricator via llvm-commits]

hanchung updated this revision to Diff 245746.
hanchung added a comment.

Load/Store the element type directly, and check the indices for AccessChain.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D74874/new/

https://reviews.llvm.org/D74874

Files:
  mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
  mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir


Index: mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir
===================================================================
--- mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir
+++ mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir
@@ -312,3 +312,41 @@
 func @memref_type(%arg0: memref<3xi1>) {
   return
 }
+
+// CHECK-LABEL: @load_store_zero_rank_float
+// CHECK: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x f32 [4]> [0]>, StorageBuffer>,
+// CHECK: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x f32 [4]> [0]>, StorageBuffer>)
+func @load_store_zero_rank_float(%arg0: memref<f32>, %arg1: memref<f32>) {
+  //      CHECK: [[ZERO1:%.*]] = spv.constant 0 : i32
+  //      CHECK: spv.AccessChain [[ARG0]][
+  // CHECK-SAME: [[ZERO1]], [[ZERO1]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Load "StorageBuffer" %{{.*}} : f32
+  %0 = load %arg0[] : memref<f32>
+  //      CHECK: [[ZERO2:%.*]] = spv.constant 0 : i32
+  //      CHECK: spv.AccessChain [[ARG1]][
+  // CHECK-SAME: [[ZERO2]], [[ZERO2]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Store "StorageBuffer" %{{.*}} : f32
+  store %0, %arg1[] : memref<f32>
+  return
+}
+
+// CHECK-LABEL: @load_store_zero_rank_int
+// CHECK: [[ARG0:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x i32 [4]> [0]>, StorageBuffer>,
+// CHECK: [[ARG1:%.*]]: !spv.ptr<!spv.struct<!spv.array<1 x i32 [4]> [0]>, StorageBuffer>)
+func @load_store_zero_rank_int(%arg0: memref<i32>, %arg1: memref<i32>) {
+  //      CHECK: [[ZERO1:%.*]] = spv.constant 0 : i32
+  //      CHECK: spv.AccessChain [[ARG0]][
+  // CHECK-SAME: [[ZERO1]], [[ZERO1]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Load "StorageBuffer" %{{.*}} : i32
+  %0 = load %arg0[] : memref<i32>
+  //      CHECK: [[ZERO2:%.*]] = spv.constant 0 : i32
+  //      CHECK: spv.AccessChain [[ARG1]][
+  // CHECK-SAME: [[ZERO2]], [[ZERO2]]
+  // CHECK-SAME: ] :
+  //      CHECK: spv.Store "StorageBuffer" %{{.*}} : i32
+  store %0, %arg1[] : memref<i32>
+  return
+}
Index: mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
===================================================================
--- mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
+++ mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp
@@ -69,6 +69,9 @@
     if (!elementSize) {
       return llvm::None;
     }
+    if (memRefType.getRank() == 0) {
+      return elementSize;
+    }
     auto dims = memRefType.getShape();
     if (llvm::is_contained(dims, ShapedType::kDynamicSize) ||
         offset == MemRefType::getDynamicStrideOrOffset() ||
@@ -332,8 +335,12 @@
   }
   SmallVector<Value, 2> linearizedIndices;
   // Add a '0' at the start to index into the struct.
-  linearizedIndices.push_back(builder.create<spirv::ConstantOp>(
-      loc, indexType, IntegerAttr::get(indexType, 0)));
+  auto zero = spirv::ConstantOp::getZero(indexType, loc, &builder);
+  linearizedIndices.push_back(zero);
+  // If it is a zero-rank memref type, extract the element directly.
+  if (!ptrLoc) {
+    ptrLoc = zero;
+  }
   linearizedIndices.push_back(ptrLoc);
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D74874.245746.patch
Type: text/x-patch
Size: 3080 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200220/ead8528c/attachment.bin>