[mlir][memref][spirv] Add SPIR-V Image Lowering (#150978)

Adds an initial conversion in the Memref -> SPIR-V lowering for images.
Any memref in the "Image" storage-class/address-space will be considered
for lowering to the `!spirv.image` type during Memref to SPIR-V
conversion. Initially only the reading of sampled images are support and
images are read via the `OpImageFetch` instruction. Future work should
expand the conversion patterns to target non-sampled images and add
support for image write operations.

Images are supported for fp32, fp16, int32, uint32, int16 and uint16
types and lit tests have been added to verify this is the case along
with negative testing to check the cases where images aren't supported.

---------

Signed-off-by: Jack Frankland <[email protected]>

Author: FranklandJack

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

@@ -244,6 +244,16 @@ class LoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts memref.load to spirv.Image + spirv.ImageFetch
+class ImageLoadOpPattern final : public OpConversionPattern<memref::LoadOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts memref.store to spirv.Store on integers.
 class IntStoreOpPattern final : public OpConversionPattern<memref::StoreOp> {
 public:
@@ -528,6 +538,17 @@ IntLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   if (!memrefType.getElementType().isSignlessInteger())
     return failure();
 
+  auto memorySpaceAttr =
+      dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
+  if (!memorySpaceAttr)
+    return rewriter.notifyMatchFailure(
+        loadOp, "missing memory space SPIR-V storage class attribute");
+
+  if (memorySpaceAttr.getValue() == spirv::StorageClass::Image)
+    return rewriter.notifyMatchFailure(
+        loadOp,
+        "failed to lower memref in image storage class to storage buffer");
+
   const auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
   Value accessChain =
       spirv::getElementPtr(typeConverter, memrefType, adaptor.getMemref(),
@@ -644,6 +665,18 @@ LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   auto memrefType = cast<MemRefType>(loadOp.getMemref().getType());
   if (memrefType.getElementType().isSignlessInteger())
     return failure();
+
+  auto memorySpaceAttr =
+      dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
+  if (!memorySpaceAttr)
+    return rewriter.notifyMatchFailure(
+        loadOp, "missing memory space SPIR-V storage class attribute");
+
+  if (memorySpaceAttr.getValue() == spirv::StorageClass::Image)
+    return rewriter.notifyMatchFailure(
+        loadOp,
+        "failed to lower memref in image storage class to storage buffer");
+
   Value loadPtr = spirv::getElementPtr(
       *getTypeConverter<SPIRVTypeConverter>(), memrefType, adaptor.getMemref(),
       adaptor.getIndices(), loadOp.getLoc(), rewriter);
@@ -662,6 +695,87 @@ LoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   return success();
 }
 
+LogicalResult
+ImageLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
+                                    ConversionPatternRewriter &rewriter) const {
+  auto memrefType = cast<MemRefType>(loadOp.getMemref().getType());
+
+  auto memorySpaceAttr =
+      dyn_cast_if_present<spirv::StorageClassAttr>(memrefType.getMemorySpace());
+  if (!memorySpaceAttr)
+    return rewriter.notifyMatchFailure(
+        loadOp, "missing memory space SPIR-V storage class attribute");
+
+  if (memorySpaceAttr.getValue() != spirv::StorageClass::Image)
+    return rewriter.notifyMatchFailure(
+        loadOp, "failed to lower memref in non-image storage class to image");
+
+  Value loadPtr = adaptor.getMemref();
+  auto memoryRequirements = calculateMemoryRequirements(loadPtr, loadOp);
+  if (failed(memoryRequirements))
+    return rewriter.notifyMatchFailure(
+        loadOp, "failed to determine memory requirements");
+
+  const auto [memoryAccess, alignment] = *memoryRequirements;
+
+  if (!loadOp.getMemRefType().hasRank())
+    return rewriter.notifyMatchFailure(
+        loadOp, "cannot lower unranked memrefs to SPIR-V images");
+
+  // We currently only support lowering of scalar memref elements to texels in
+  // the R[16|32][f|i|ui] formats. Future work will enable lowering of vector
+  // elements to texels in richer formats.
+  if (!isa<spirv::ScalarType>(loadOp.getMemRefType().getElementType()))
+    return rewriter.notifyMatchFailure(
+        loadOp,
+        "cannot lower memrefs who's element type is not a SPIR-V scalar type"
+        "to SPIR-V images");
+
+  // We currently only support sampled images since OpImageFetch does not work
+  // for plain images and the OpImageRead instruction needs to be materialized
+  // instead or texels need to be accessed via atomics through a texel pointer.
+  // Future work will generalize support to plain images.
+  auto convertedPointeeType = cast<spirv::PointerType>(
+      getTypeConverter()->convertType(loadOp.getMemRefType()));
+  if (!isa<spirv::SampledImageType>(convertedPointeeType.getPointeeType()))
+    return rewriter.notifyMatchFailure(loadOp,
+                                       "cannot lower memrefs which do not "
+                                       "convert to SPIR-V sampled images");
+
+  // Materialize the lowering.
+  Location loc = loadOp->getLoc();
+  auto imageLoadOp =
+      spirv::LoadOp::create(rewriter, loc, loadPtr, memoryAccess, alignment);
+  // Extract the image from the sampled image.
+  auto imageOp = spirv::ImageOp::create(rewriter, loc, imageLoadOp);
+
+  // Build a vector of coordinates or just a scalar index if we have a 1D image.
+  Value coords;
+  if (memrefType.getRank() != 1) {
+    auto coordVectorType = VectorType::get({loadOp.getMemRefType().getRank()},
+                                           adaptor.getIndices().getType()[0]);
+    coords = spirv::CompositeConstructOp::create(rewriter, loc, coordVectorType,
+                                                 adaptor.getIndices());
+  } else {
+    coords = adaptor.getIndices()[0];
+  }
+
+  // Fetch the value out of the image.
+  auto resultVectorType = VectorType::get({4}, loadOp.getType());
+  auto fetchOp = spirv::ImageFetchOp::create(
+      rewriter, loc, resultVectorType, imageOp, coords,
+      mlir::spirv::ImageOperandsAttr{}, ValueRange{});
+
+  // Note that because OpImageFetch returns a rank 4 vector we need to extract
+  // the elements corresponding to the load which will since we only support the
+  // R[16|32][f|i|ui] formats will always be the R(red) 0th vector element.
+  auto compositeExtractOp =
+      spirv::CompositeExtractOp::create(rewriter, loc, fetchOp, 0);
+
+  rewriter.replaceOp(loadOp, compositeExtractOp);
+  return success();
+}
+
 LogicalResult
 IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
@@ -953,11 +1067,11 @@ LogicalResult ExtractAlignedPointerAsIndexOpPattern::matchAndRewrite(
 namespace mlir {
 void populateMemRefToSPIRVPatterns(const SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
-  patterns
-      .add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
-           DeallocOpPattern, IntLoadOpPattern, IntStoreOpPattern, LoadOpPattern,
-           MemorySpaceCastOpPattern, StoreOpPattern, ReinterpretCastPattern,
-           CastPattern, ExtractAlignedPointerAsIndexOpPattern>(
-          typeConverter, patterns.getContext());
+  patterns.add<AllocaOpPattern, AllocOpPattern, AtomicRMWOpPattern,
+               DeallocOpPattern, IntLoadOpPattern, ImageLoadOpPattern,
+               IntStoreOpPattern, LoadOpPattern, MemorySpaceCastOpPattern,
+               StoreOpPattern, ReinterpretCastPattern, CastPattern,
+               ExtractAlignedPointerAsIndexOpPattern>(typeConverter,
+                                                      patterns.getContext());
 }
 } // namespace mlir

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

@@ -608,6 +608,45 @@ static Type convertSubByteMemrefType(const spirv::TargetEnv &targetEnv,
   return wrapInStructAndGetPointer(arrayType, storageClass);
 }
 
+static spirv::Dim convertRank(int64_t rank) {
+  switch (rank) {
+  case 1:
+    return spirv::Dim::Dim1D;
+  case 2:
+    return spirv::Dim::Dim2D;
+  case 3:
+    return spirv::Dim::Dim3D;
+  default:
+    llvm_unreachable("Invalid memref rank!");
+  }
+}
+
+static spirv::ImageFormat getImageFormat(Type elementType) {
+  return llvm::TypeSwitch<Type, spirv::ImageFormat>(elementType)
+      .Case<Float16Type>([](Float16Type) { return spirv::ImageFormat::R16f; })
+      .Case<Float32Type>([](Float32Type) { return spirv::ImageFormat::R32f; })
+      .Case<IntegerType>([](IntegerType intType) {
+        auto const isSigned = intType.isSigned() || intType.isSignless();
+#define BIT_WIDTH_CASE(BIT_WIDTH)                                              \
+  case BIT_WIDTH:                                                              \
+    return isSigned ? spirv::ImageFormat::R##BIT_WIDTH##i                      \
+                    : spirv::ImageFormat::R##BIT_WIDTH##ui
+
+        switch (intType.getWidth()) {
+          BIT_WIDTH_CASE(16);
+          BIT_WIDTH_CASE(32);
+        default:
+          llvm_unreachable("Unhandled integer type!");
+        }
+      })
+      .Default([](Type) {
+        llvm_unreachable("Unhandled element type!");
+        // We need to return something here to satisfy the type switch.
+        return spirv::ImageFormat::R32f;
+      });
+#undef BIT_WIDTH_CASE
+}
+
 static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
                               const SPIRVConversionOptions &options,
                               MemRefType type) {
@@ -623,6 +662,41 @@ static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
   }
   spirv::StorageClass storageClass = attr.getValue();
 
+  // Images are a special case since they are an opaque type from which elements
+  // may be accessed via image specific ops or directly through a texture
+  // pointer.
+  if (storageClass == spirv::StorageClass::Image) {
+    const int64_t rank = type.getRank();
+    if (rank < 1 || rank > 3) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << type << " illegal: cannot lower memref of rank " << rank
+                 << " to a SPIR-V Image\n");
+      return nullptr;
+    }
+
+    // Note that we currently only support lowering to single element texels
+    // e.g. R32f.
+    auto elementType = type.getElementType();
+    if (!isa<spirv::ScalarType>(elementType)) {
+      LLVM_DEBUG(llvm::dbgs() << type << " illegal: cannot lower memref of "
+                              << elementType << " to a  SPIR-V Image\n");
+      return nullptr;
+    }
+
+    // Currently every memref in the image storage class is converted to a
+    // sampled image so we can hardcode the NeedSampler field. Future work
+    // will generalize this to support regular non-sampled images.
+    auto spvImageType = spirv::ImageType::get(
+        elementType, convertRank(rank), spirv::ImageDepthInfo::DepthUnknown,
+        spirv::ImageArrayedInfo::NonArrayed,
+        spirv::ImageSamplingInfo::SingleSampled,
+        spirv::ImageSamplerUseInfo::NeedSampler, getImageFormat(elementType));
+    auto spvSampledImageType = spirv::SampledImageType::get(spvImageType);
+    auto imagePtrType = spirv::PointerType::get(
+        spvSampledImageType, spirv::StorageClass::UniformConstant);
+    return imagePtrType;
+  }
+
   if (isa<IntegerType>(type.getElementType())) {
     if (type.getElementTypeBitWidth() == 1)
       return convertBoolMemrefType(targetEnv, options, type, storageClass);