[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 created this revision.
hanchung added reviewers: mravishankar, antiagainst.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a project: LLVM.

Repository:
  rG LLVM Github Monorepo

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,25 @@
 func @memref_type(%arg0: memref<3xi1>) {
   return
 }
+
+// CHECK-LABEL: @load_store_zero_rank_float
+// CHECK: !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: spv.AccessChain
+  // CHECK: spv.Load
+  // CHECK: spv.Store
+  %0 = load %arg0[] : memref<f32>
+  store %0, %arg1[] : memref<f32>
+  return
+}
+
+// CHECK-LABEL: @load_store_zero_rank_int
+// CHECK: !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: spv.AccessChain
+  // CHECK: spv.Load
+  // CHECK: spv.Store
+  %0 = load %arg0[] : memref<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() ||
@@ -334,7 +337,9 @@
   // Add a '0' at the start to index into the struct.
   linearizedIndices.push_back(builder.create<spirv::ConstantOp>(
       loc, indexType, IntegerAttr::get(indexType, 0)));
-  linearizedIndices.push_back(ptrLoc);
+  if (ptrLoc) {
+    linearizedIndices.push_back(ptrLoc);
+  }
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D74874.245535.patch
Type: text/x-patch
Size: 1940 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200219/1bf514cd/attachment.bin>