[mlir][spirv] Allow bitwidth emulation on runtime arrays

[antiagainst]

Summary

Runtime arrays are converted from memrefs with unknown
dimensions.

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

998

Nit: Having dstType initially being the StructType and overwriting it with element type adds cognitive overhead to me while reading. Can we just use elementType for the actual element type.

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp,
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, loadOperands.memref(),
                            loadOperands.indices(), loc, rewriter);
 
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
-  auto dstType = typeConverter.convertType(memrefType)
+  Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
-                     .getPointeeType()
-                     .cast<spirv::StructType>()
-                     .getElementType(0)
-                     .cast<spirv::ArrayType>()
-                     .getElementType();
+                         .getPointeeType();
+  Type structElemType = pointeeType.cast<spirv::StructType>().getElementType(0);
+  Type dstType;
+  if (auto arrayType = structElemType.dyn_cast<spirv::ArrayType>())
+    dstType = arrayType.getElementType();
+  else
+    dstType = structElemType.cast<spirv::RuntimeArrayType>().getElementType();
+
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
   // If the rewrited load op has the same bit width, use the loading value
   // directly.
   if (srcBits == dstBits) {
     rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp,
                                                accessChainOp.getResult());
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp,
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
                            storeOperands.indices(), loc, rewriter);
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
 
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
-  auto dstType = typeConverter.convertType(memrefType)
+  Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
-                     .getPointeeType()
-                     .cast<spirv::StructType>()
-                     .getElementType(0)
-                     .cast<spirv::ArrayType>()
-                     .getElementType();
+                         .getPointeeType();
+  Type structElemType = pointeeType.cast<spirv::StructType>().getElementType(0);
+  Type dstType;
+  if (auto arrayType = structElemType.dyn_cast<spirv::ArrayType>())
+    dstType = arrayType.getElementType();
+  else
+    dstType = structElemType.cast<spirv::RuntimeArrayType>().getElementType();
+
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
   if (srcBits == dstBits) {
     rewriter.replaceOpWithNewOp<spirv::StoreOp>(
         storeOp, accessChainOp.getResult(), storeOperands.value());
     return success();
   }

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

func @load_store_zero_rank_int(%arg0: memref<i32>, %arg1: memref<i32>) {
   //      CHECK: spv.AccessChain [[ARG1]][
   // CHECK-SAME: [[ZERO2]], [[ZERO2]]
   // CHECK-SAME: ] :
   //      CHECK: spv.Store "StorageBuffer" %{{.*}} : i32
   memref.store %0, %arg1[] : memref<i32>
   return
 }
 
+// CHECK-LABEL: func @load_store_unknown_dim
+// CHECK-SAME: %[[SRC:[a-z0-9]+]]: !spv.ptr<!spv.struct<(!spv.rtarray<i32, stride=4> [0])>, StorageBuffer>,
+// CHECK-SAME: %[[DST:[a-z0-9]+]]: !spv.ptr<!spv.struct<(!spv.rtarray<i32, stride=4> [0])>, StorageBuffer>)
+func @load_store_unknown_dim(%i: index, %source: memref<?xi32>, %dest: memref<?xi32>) {
+  // CHECK: %[[AC0:.+]] = spv.AccessChain %[[SRC]]
+  // CHECK: spv.Load "StorageBuffer" %[[AC0]]
+  %0 = memref.load %source[%i] : memref<?xi32>
+  // CHECK: %[[AC1:.+]] = spv.AccessChain %[[DST]]
+  // CHECK: spv.Store "StorageBuffer" %[[AC1]]
+  memref.store %0, %dest[%i]: memref<?xi32>
+  return
+}
+
 } // end module
 
 // -----
 
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
 // TODO: Test i64 types.
 module attributes {