[mlir][spirv] Handle zero-sized memrefs

[kuhar]

Summary

Make sure to check type conversion results. Add missing tests.

Fix some typos in the surrounding code.

Fixes: https://github.com/llvm/llvm-project/issues/64409

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = allocaOp.getType();
   if (!isAllocationSupported(allocaOp, allocType))
     return rewriter.notifyMatchFailure(allocaOp, "unhandled allocation type");
 
   // Get the SPIR-V type for the allocation.
   Type spirvType = getTypeConverter()->convertType(allocType);
+  if (!spirvType)
+    return rewriter.notifyMatchFailure(allocaOp, "type conversion failed");
+
   rewriter.replaceOpWithNewOp<spirv::VariableOp>(allocaOp, spirvType,
                                                  spirv::StorageClass::Function,
                                                  /*initializer=*/nullptr);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 AllocOpPattern::matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = operation.getType();
   if (!isAllocationSupported(operation, allocType))
     return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
 
   // Get the SPIR-V type for the allocation.
   Type spirvType = getTypeConverter()->convertType(allocType);
+  if (!spirvType)
+    return rewriter.notifyMatchFailure(operation, "type conversion failed");
 
   // Insert spirv.GlobalVariable for this allocation.
   Operation *parent =
       SymbolTable::getNearestSymbolTable(operation->getParentOp());
   if (!parent)
     return failure();
   Location loc = operation.getLoc();
   spirv::GlobalVariableOp varOp;
if (typeConverter.allows(spirv::Capability::Kernel)) {
     if (auto arrayType = dyn_cast<spirv::ArrayType>(structElemType))
       dstType = arrayType.getElementType();
     else
       dstType = cast<spirv::RuntimeArrayType>(structElemType).getElementType();
   }
   int dstBits = dstType.getIntOrFloatBitWidth();
   assert(dstBits % srcBits == 0);
 
-  // If the rewrited load op has the same bit width, use the loading value
+  // If the rewritten load op has the same bit width, use the loading value
   // directly.
   if (srcBits == dstBits) {
     Value loadVal = rewriter.create<spirv::LoadOp>(loc, accessChain);
     if (isBool)
       loadVal = castIntNToBool(loc, loadVal, rewriter);
     rewriter.replaceOp(loadOp, loadVal);
     return success();
   }
return rewriter.notifyMatchFailure(loc, [resultType](Diagnostic &diag) {
            << " must be a SPIR-V storage class";
     });
 
   spirv::StorageClass sourceSc = sourceStorageClassAttr.getValue();
   spirv::StorageClass resultSc = resultStorageClassAttr.getValue();
 
   Value result = adaptor.getSource();
   Type resultPtrType = typeConverter.convertType(resultType);
+  if (!resultPtrType)
+    return rewriter.notifyMatchFailure(addrCastOp,
+                                       "failed to convert memref type");
+
   Type genericPtrType = resultPtrType;
   // SPIR-V doesn't have a general address space cast operation. Instead, it has
   // conversions to and from generic pointers. To implement the general case,
   // we use specific-to-generic conversions when the source class is not
   // generic. Then when the result storage class is not generic, we convert the
-  // generic pointer (either the input on ar intermediate result) to theat
+  // generic pointer (either the input on ar intermediate result) to that
   // class. This also means that we'll need the intermediate generic pointer
   // type if neither the source or destination have it.
   if (sourceSc != spirv::StorageClass::Generic &&
       resultSc != spirv::StorageClass::Generic) {
     Type intermediateType =
         MemRefType::get(sourceType.getShape(), sourceType.getElementType(),
                         sourceType.getLayout(),
                         rewriter.getAttr<spirv::StorageClassAttr>(

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
 
   std::optional<int64_t> memrefSize = getTypeNumBytes(options, type);
   if (!memrefSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce element count\n");
     return nullptr;
   }
 
+  if (*memrefSize == 0) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: zero-element memrefs are not supported\n");
+    return nullptr;
+  }
+
   int64_t arrayElemCount = llvm::divideCeil(*memrefSize, *arrayElemSize);
   int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
   auto arrayType = spirv::ArrayType::get(arrayElemType, arrayElemCount, stride);
   if (targetEnv.allows(spirv::Capability::Kernel))
     return spirv::PointerType::get(arrayType, storageClass);
   return wrapInStructAndGetPointer(arrayType, storageClass);
 }
 

mlir/test/Conversion/MemRefToSPIRV/alloc.mlir

 //       CHECK:   %[[LOC:.+]] = spirv.SDiv
 //       CHECK:   %[[PTR:.+]] = spirv.AccessChain %[[VAR]][%{{.+}}, %[[LOC]]]
 //       CHECK:   %{{.+}} = spirv.Load "Workgroup" %[[PTR]] : i32
 //       CHECK:   %[[LOC:.+]] = spirv.SDiv
 //       CHECK:   %[[PTR:.+]] = spirv.AccessChain %[[VAR]][%{{.+}}, %[[LOC]]]
 //       CHECK:   %{{.+}} = spirv.AtomicAnd "Workgroup" "AcquireRelease" %[[PTR]], %{{.+}} : !spirv.ptr<i32, Workgroup>
 //       CHECK:   %{{.+}} = spirv.AtomicOr "Workgroup" "AcquireRelease" %[[PTR]], %{{.+}} : !spirv.ptr<i32, Workgroup>
 
-
 // -----
 
 module attributes {
   spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
   }
 {
   func.func @two_allocs() {

 }
 
 //   CHECK-DAG: spirv.GlobalVariable @__workgroup_mem__{{[0-9]+}}
 //  CHECK-SAME:   !spirv.ptr<!spirv.struct<(!spirv.array<2 x vector<2xi32>>)>, Workgroup>
 //   CHECK-DAG: spirv.GlobalVariable @__workgroup_mem__{{[0-9]+}}
 //  CHECK-SAME:   !spirv.ptr<!spirv.struct<(!spirv.array<4 x vector<4xf32>>)>, Workgroup>
 // CHECK-LABEL: func @two_allocs_vector()
 
-
 // -----
 
 module attributes {
   spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
   }
 {
   // CHECK-LABEL: func @alloc_dynamic_size

 // CHECK-LABEL: func @alloc_dealloc_workgroup_mem_kernel
 //   CHECK-NOT:   memref.alloc
 //       CHECK:   %[[PTR:.+]] = spirv.mlir.addressof @[[$VAR]]
 //       CHECK:   %[[LOADPTR:.+]] = spirv.AccessChain %[[PTR]]
 //       CHECK:   %[[VAL:.+]] = spirv.Load "Workgroup" %[[LOADPTR]] : f32
 //       CHECK:   %[[STOREPTR:.+]] = spirv.AccessChain %[[PTR]]
 //       CHECK:   spirv.Store "Workgroup" %[[STOREPTR]], %[[VAL]] : f32
 //   CHECK-NOT:   memref.dealloc
+
+// -----
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+  }
+{
+  func.func @zero_size() {
+    %0 = memref.alloc() : memref<0xf32, #spirv.storage_class<Workgroup>>
+    return
+  }
+}
+
+// Zero-sized allocations are not handled yet. Just make sure we do not crash.
+// CHECK-LABEL: func @zero_size()

mlir/test/Conversion/MemRefToSPIRV/alloca.mlir

 
 // CHECK-LABEL: func @alloc_function_variable
 //       CHECK:   %[[VAR:.+]] = spirv.Variable : !spirv.ptr<!spirv.struct<(!spirv.array<20 x f32>)>, Function>
 //       CHECK:   %[[LOADPTR:.+]] = spirv.AccessChain %[[VAR]]
 //       CHECK:   %[[VAL:.+]] = spirv.Load "Function" %[[LOADPTR]] : f32
 //       CHECK:   %[[STOREPTR:.+]] = spirv.AccessChain %[[VAR]]
 //       CHECK:   spirv.Store "Function" %[[STOREPTR]], %[[VAL]] : f32
 
-
 // -----
 
 module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
   func.func @two_allocs() {
     %0 = memref.alloca() : memref<4x5xf32, #spirv.storage_class<Function>>
     %1 = memref.alloca() : memref<2x3xi32, #spirv.storage_class<Function>>
     return
   }
module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
   // CHECK-LABEL: func @alloc_unsupported_memory_space
   func.func @alloc_unsupported_memory_space(%arg0: index) -> f32 {
     // CHECK: memref.alloca
     %0 = memref.alloca() : memref<4x5xf32>
     %1 = memref.load %0[%arg0, %arg0] : memref<4x5xf32>
     return %1: f32
   }
 }
+
+// -----
+
+module attributes {spirv.target_env = #spirv.target_env<#spirv.vce<v1.3, [Shader], []>, #spirv.resource_limits<>>} {
+  func.func @zero_size() {
+    %0 = memref.alloca() : memref<0xf32, #spirv.storage_class<Function>>
+    return
+  }
+}
+
+// Zero-sized allocations are not handled yet. Just make sure we do not crash.
+// CHECK-LABEL: func @zero_size

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

 //  CHECK-SAME:  (%[[MEM:.*]]: memref<4x?xf32, #spirv.storage_class<CrossWorkgroup>>)
 func.func @cast_to_static(%arg: memref<4x?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<4x4xf32, #spirv.storage_class<CrossWorkgroup>> {
 //       CHECK:  %[[MEM1:.*]] =  memref.cast %[[MEM]] : memref<4x?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<4x4xf32, #spirv.storage_class<CrossWorkgroup>>
 //       CHECK:  return %[[MEM1]]
   %ret = memref.cast %arg : memref<4x?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<4x4xf32, #spirv.storage_class<CrossWorkgroup>>
   return %ret : memref<4x4xf32, #spirv.storage_class<CrossWorkgroup>>
 }
 
+// TODO: Not supported yet
+// CHECK-LABEL: func.func @cast_to_static_zero_elems
+//  CHECK-SAME:  (%[[MEM:.*]]: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>)
+func.func @cast_to_static_zero_elems(%arg: memref<?xf32, #spirv.storage_class<CrossWorkgroup>>) -> memref<0xf32, #spirv.storage_class<CrossWorkgroup>> {
+//       CHECK:  %[[MEM1:.*]] =  memref.cast %[[MEM]] : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<0xf32, #spirv.storage_class<CrossWorkgroup>>
+//       CHECK:  return %[[MEM1]]
+  %ret = memref.cast %arg : memref<?xf32, #spirv.storage_class<CrossWorkgroup>> to memref<0xf32, #spirv.storage_class<CrossWorkgroup>>
+  return %ret : memref<0xf32, #spirv.storage_class<CrossWorkgroup>>
+}
+
 }