[mlir][memref][spirv]: Address Feedback

* Outline lambdas to helper functions.
* Removed incorrect `const` usage.
* Generalize error conditions to SPIR-V scalar types.
* Create local variable for debug location.

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

Author: FranklandJack

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

@@ -724,10 +724,11 @@ ImageLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
   // 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 (!loadOp.getMemRefType().getElementType().isIntOrFloat())
+  if (!isa<spirv::ScalarType>(loadOp.getMemRefType().getElementType()))
     return rewriter.notifyMatchFailure(
-        loadOp, "cannot lower memrefs who's element type is not int or float "
-                "to SPIR-V images");
+        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
@@ -741,34 +742,34 @@ ImageLoadOpPattern::matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                                        "convert to SPIR-V sampled images");
 
   // Materialize the lowering.
-  auto imageLoadOp = spirv::LoadOp::create(rewriter, loadOp->getLoc(), loadPtr,
-                                           memoryAccess, alignment);
+  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, loadOp->getLoc(), imageLoadOp);
+  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) {
-    const auto coordVectorType = VectorType::get(
-        {loadOp.getMemRefType().getRank()}, adaptor.getIndices().getType()[0]);
-    coords = spirv::CompositeConstructOp::create(
-        rewriter, loadOp->getLoc(), coordVectorType, adaptor.getIndices());
+    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, loadOp->getLoc(), resultVectorType, imageOp, coords,
+      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, loadOp->getLoc(), fetchOp, 0);
+      spirv::CompositeExtractOp::create(rewriter, loc, fetchOp, 0);
 
   rewriter.replaceOp(loadOp, compositeExtractOp);
   return success();

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

@@ -560,6 +560,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) {
@@ -587,64 +626,23 @@ static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
       return nullptr;
     }
 
-    const auto dim = [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!");
-      }
-    }();
-
     // Note that we currently only support lowering to single element texels
     // e.g. R32f.
     auto elementType = type.getElementType();
-    if (!elementType.isIntOrFloat()) {
+    if (!isa<spirv::ScalarType>(elementType)) {
       LLVM_DEBUG(llvm::dbgs() << type << " illegal: cannot lower memref of "
                               << elementType << " to a  SPIR-V Image\n");
       return nullptr;
     }
 
-    const auto imageFormat = [&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
-    }();
-
     // 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, dim, spirv::ImageDepthInfo::DepthUnknown,
+        elementType, convertRank(rank), spirv::ImageDepthInfo::DepthUnknown,
         spirv::ImageArrayedInfo::NonArrayed,
         spirv::ImageSamplingInfo::SingleSampled,
-        spirv::ImageSamplerUseInfo::NeedSampler, imageFormat);
+        spirv::ImageSamplerUseInfo::NeedSampler, getImageFormat(elementType));
     auto spvSampledImageType = spirv::SampledImageType::get(spvImageType);
     auto imagePtrType = spirv::PointerType::get(
         spvSampledImageType, spirv::StorageClass::UniformConstant);