[mlir][spirv] Fix crash in convert-gpu-to-spirv pass with memrefs with affine maps

[Hardcode84]

Summary

spirv::getElementPtr can return null (for memrefs with affine map) but patterns didn't handle this.

[antiagainst]

Thanks for fixing this! Just one nit: could you also update the comment for spirv::getElementPtr to mention that it can return nullptr when it cannot perform the index calculation?

[Hardcode84]

Add comment to getElementPtr

[Hardcode84]

Also, is there any plans to actually support memrefs with dynamic offset and strides?

[antiagainst]

Awesome, thanks!

In D106988#2911568, @Hardcode84 wrote:

Also, is there any plans to actually support memrefs with dynamic offset and strides?

Dynamic shapes / strides are a bit tricky for the SPIR-V side due to that we cannot always use a memref descriptor for it like LLVM (see the details here). In IREE we have been working with another approach towards it, without the memref descriptor. Basically we flatten the n-D memref into 1-D (effectively what the memref descriptor is doing under the hood for CPU side, relying on LLVM proper) before converting to SPIR-V using this pass. And then SPIR-V conversion only needs to handle the case where we have a 1-D flat memref/buffer. It effectively pushes the index calculation to a previous stage; instead of relying it being done at SPIR-V level. I'd like to upstream that sometime later.

mlir/include/mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h

 
 /// Generates IR to perform index linearization with the given `indices` and
 /// their corresponding `strides`, adding an initial `offset`.
 Value linearizeIndex(ValueRange indices, ArrayRef<int64_t> strides,
                      int64_t offset, Location loc, OpBuilder &builder);
 
 /// Performs the index computation to get to the element at `indices` of the
 /// memory pointed to by `basePtr`, using the layout map of `baseType`.
+/// Returns null if index computation cannot be performed.
 
 // TODO: This method assumes that the `baseType` is a MemRefType with AffineMap
 // that has static strides. Extend to handle dynamic strides.
 spirv::AccessChainOp getElementPtr(SPIRVTypeConverter &typeConverter,
                                    MemRefType baseType, Value basePtr,
                                    ValueRange indices, Location loc,
                                    OpBuilder &builder);
 
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_SPIRVCONVERSION_H

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp,
   if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, loadOperands.memref(),
                            loadOperands.indices(), loc, rewriter);
 
+  if (!accessChainOp)
+    return failure();
+
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
   Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
                          .getPointeeType();
   Type structElemType = pointeeType.cast<spirv::StructType>().getElementType(0);
LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, ArrayRef<Value> operands,
                                ConversionPatternRewriter &rewriter) const {
   memref::LoadOpAdaptor loadOperands(operands);
   auto memrefType = loadOp.memref().getType().cast<MemRefType>();
   if (memrefType.getElementType().isSignlessInteger())
     return failure();
   auto loadPtr = spirv::getElementPtr(
       *getTypeConverter<SPIRVTypeConverter>(), memrefType,
       loadOperands.memref(), loadOperands.indices(), loadOp.getLoc(), rewriter);
+
+  if (!loadPtr)
+    return failure();
+
   rewriter.replaceOpWithNewOp<spirv::LoadOp>(loadOp, loadPtr);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp,
   if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
   auto loc = storeOp.getLoc();
   auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   spirv::AccessChainOp accessChainOp =
       spirv::getElementPtr(typeConverter, memrefType, storeOperands.memref(),
                            storeOperands.indices(), loc, rewriter);
+
+  if (!accessChainOp)
+    return failure();
+
   int srcBits = memrefType.getElementType().getIntOrFloatBitWidth();
 
   bool isBool = srcBits == 1;
   if (isBool)
     srcBits = typeConverter.getOptions().boolNumBits;
   Type pointeeType = typeConverter.convertType(memrefType)
                          .cast<spirv::PointerType>()
                          .getPointeeType();
StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp,
   memref::StoreOpAdaptor storeOperands(operands);
   auto memrefType = storeOp.memref().getType().cast<MemRefType>();
   if (memrefType.getElementType().isSignlessInteger())
     return failure();
   auto storePtr =
       spirv::getElementPtr(*getTypeConverter<SPIRVTypeConverter>(), memrefType,
                            storeOperands.memref(), storeOperands.indices(),
                            storeOp.getLoc(), rewriter);
+
+  if (!storePtr)
+    return failure();
+
   rewriter.replaceOpWithNewOp<spirv::StoreOp>(storeOp, storePtr,
                                               storeOperands.value());
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // XorOp
 //===----------------------------------------------------------------------===//