[mlir][spirv] Lower memref with dynamic dimensions to runtime arrays

[antiagainst]

Summary

memref types with dynamic dimensions do not have a compile-time
known size. They should be mapped to SPIR-V runtime array types.

[mravishankar]

This is great! Thanks Lei, The type conversion looks good, but its hard to see how this works without seeing how it works in practice without seeing it work with loads/stores.

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

334–336

Nit: Good to avoid using auto here cause the type of the result is not obvious.

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

if (!arrayElemSize) {
     return llvm::None;
   }
 
   return spirv::ArrayType::get(*arrayElemType, arrayElemCount, *arrayElemSize);
 }
 
 static Optional<Type> convertMemrefType(const spirv::TargetEnv &targetEnv,
                                         MemRefType type) {
-  // TODO(ravishankarm) : Handle dynamic shapes.
-  if (!type.hasStaticShape()) {
+  Optional<spirv::StorageClass> storageClass =
+      SPIRVTypeConverter::getStorageClassForMemorySpace(type.getMemorySpace());
+  if (!storageClass) {
     LLVM_DEBUG(llvm::dbgs()
-               << type << " illegal: dynamic shape unimplemented\n");
+               << type << " illegal: cannot convert memory space\n");
     return llvm::None;
   }
 
   auto scalarType = type.getElementType().dyn_cast<spirv::ScalarType>();
   if (!scalarType) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert non-scalar element type\n");
     return llvm::None;
   }
 
+  auto arrayElemType = convertScalarType(targetEnv, scalarType, storageClass);
+  if (!arrayElemType)
+    return llvm::None;
+
   Optional<int64_t> scalarSize = getTypeNumBytes(scalarType);
-  Optional<int64_t> memrefSize = getTypeNumBytes(type);
-  if (!scalarSize || !memrefSize) {
+  if (!scalarSize) {
     LLVM_DEBUG(llvm::dbgs()
-               << type << " illegal: cannot deduce element count\n");
+               << type << " illegal: cannot deduce element size\n");
     return llvm::None;
   }
 
-  auto arrayElemCount = *memrefSize / *scalarSize;
+  if (!type.hasStaticShape()) {
+    auto arrayType = spirv::RuntimeArrayType::get(*arrayElemType, *scalarSize);
+    // Wrap in a struct to satisfy Vulkan interface requirements.
+    auto structType = spirv::StructType::get(arrayType, 0);
+    return spirv::PointerType::get(structType, *storageClass);
+  }
 
-  auto storageClass =
-      SPIRVTypeConverter::getStorageClassForMemorySpace(type.getMemorySpace());
-  if (!storageClass) {
+  Optional<int64_t> memrefSize = getTypeNumBytes(type);
+  if (!memrefSize) {
     LLVM_DEBUG(llvm::dbgs()
-               << type << " illegal: cannot convert memory space\n");
+               << type << " illegal: cannot deduce element count\n");
     return llvm::None;
   }
 
-  auto arrayElemType = convertScalarType(targetEnv, scalarType, storageClass);
-  if (!arrayElemType)
-    return llvm::None;
+  auto arrayElemCount = *memrefSize / *scalarSize;
+
   Optional<int64_t> arrayElemSize = getTypeNumBytes(*arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return llvm::None;
   }
 
   auto arrayType =

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

// CHECK-SAME: !spv.array<88 x f16, stride=2> [0]>, StorageBuffer>
   %arg8: memref<16x4xf16, offset: 0, strides: [1, 16]>,
   %arg9: memref<16x4xf16, offset: 0, strides: [1, 22]>
 ) { return }
 
 } // end module
 
 // -----
 
-// Check that dynamic shapes are not supported.
+// Dynamic shapes
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [], []>,
     {max_compute_workgroup_invocations = 128 : i32,
      max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
 } {
 
+// Check that unranked shapes are not supported.
 // CHECK-LABEL: func @unranked_memref
 // CHECK-SAME: memref<*xi32>
 func @unranked_memref(%arg0: memref<*xi32>) { return }
 
 // CHECK-LABEL: func @dynamic_dim_memref
-// CHECK-SAME: memref<8x?xi32>
-func @dynamic_dim_memref(%arg0: memref<8x?xi32>) { return }
+// CHECK-SAME: !spv.ptr<!spv.struct<!spv.rtarray<i32, stride=4> [0]>, StorageBuffer>
+// CHECK-SAME: !spv.ptr<!spv.struct<!spv.rtarray<f32, stride=4> [0]>, StorageBuffer>
+func @dynamic_dim_memref(%arg0: memref<8x?xi32>,
+                         %arg1: memref<?x?xf32>)
+{ return }
 
 } // end module
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Tensor types
 //===----------------------------------------------------------------------===//