[mlir][spirv] Support shaped types with index element

[kuhar]

Summary

This makes the SPIR-V type converter first convert index element types
to the right integer type.

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

[antiagainst]

Nice, thanks for fixing this!

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

 
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #include <functional>
 
 #define DEBUG_TYPE "mlir-spirv-conversion"
 
 using namespace mlir;
auto structType = needsExplicitLayout(storageClass)
                         : spirv::StructType::get(elementType);
   return spirv::PointerType::get(structType, storageClass);
 }
 
 //===----------------------------------------------------------------------===//
 // Type Conversion
 //===----------------------------------------------------------------------===//
 
+static spirv::ScalarType getIndexType(MLIRContext *ctx,
+                                      const SPIRVConversionOptions &options) {
+  return cast<spirv::ScalarType>(
+      IntegerType::get(ctx, options.use64bitIndex ? 64 : 32));
+}
+
 Type SPIRVTypeConverter::getIndexType() const {
-  return IntegerType::get(getContext(), options.use64bitIndex ? 64 : 32);
+  return ::getIndexType(getContext(), options);
 }
 
 MLIRContext *SPIRVTypeConverter::getContext() const {
   return targetEnv.getAttr().getContext();
 }
 
 bool SPIRVTypeConverter::allows(spirv::Capability capability) {
   return targetEnv.allows(capability);
convertScalarType(const spirv::TargetEnv &targetEnv,
   }
 
   auto intType = type.cast<IntegerType>();
   LLVM_DEBUG(llvm::dbgs() << type << " converted to 32-bit for SPIR-V\n");
   return IntegerType::get(targetEnv.getContext(), /*width=*/32,
                           intType.getSignedness());
 }
 
+/// Returns a type with the same shape but with any index element type converted
+/// to the matching integer type. This is a noop when the element type is not
+/// the index type.
+static ShapedType
+convertIndexElementType(ShapedType type,
+                        const SPIRVConversionOptions &options) {
+  Type indexType = dyn_cast<IndexType>(type.getElementType());
+  if (!indexType)
+    return type;
+
+  return type.clone(getIndexType(type.getContext(), options));
+}
+
 /// Converts a vector `type` to a suitable type under the given `targetEnv`.
 static Type
 convertVectorType(const spirv::TargetEnv &targetEnv,
                   const SPIRVConversionOptions &options, VectorType type,
                   std::optional<spirv::StorageClass> storageClass = {}) {
-  auto scalarType = type.getElementType().cast<spirv::ScalarType>();
+  type = cast<VectorType>(convertIndexElementType(type, options));
+  auto scalarType = dyn_cast_or_null<spirv::ScalarType>(type.getElementType());
+  if (!scalarType) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot convert non-scalar element type\n");
+    return nullptr;
+  }
+
   if (type.getRank() <= 1 && type.getNumElements() == 1)
     return convertScalarType(targetEnv, options, scalarType, storageClass);
 
   if (!spirv::CompositeType::isValid(type)) {
     LLVM_DEBUG(llvm::dbgs() << type << " illegal: > 4-element unimplemented\n");
     return nullptr;
   }
 
static Type convertTensorType(const spirv::TargetEnv &targetEnv,
                               TensorType type) {
   // TODO: Handle dynamic shapes.
   if (!type.hasStaticShape()) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: dynamic shape unimplemented\n");
     return nullptr;
   }
 
-  auto scalarType = type.getElementType().dyn_cast<spirv::ScalarType>();
+  type = cast<TensorType>(convertIndexElementType(type, options));
+  auto scalarType = dyn_cast_or_null<spirv::ScalarType>(type.getElementType());
   if (!scalarType) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot convert non-scalar element type\n");
     return nullptr;
   }
 
   std::optional<int64_t> scalarSize = getTypeNumBytes(options, scalarType);
   std::optional<int64_t> tensorSize = getTypeNumBytes(options, type);
static Type convertMemrefType(const spirv::TargetEnv &targetEnv,
   Type arrayElemType;
   Type elementType = type.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>()) {
     arrayElemType =
         convertVectorType(targetEnv, options, vecType, storageClass);
   } else if (auto scalarType = elementType.dyn_cast<spirv::ScalarType>()) {
     arrayElemType =
         convertScalarType(targetEnv, options, scalarType, storageClass);
+  } else if (auto indexType = elementType.dyn_cast<IndexType>()) {
+    type = convertIndexElementType(type, options).cast<MemRefType>();
+    arrayElemType = type.getElementType();
   } else {
     LLVM_DEBUG(
         llvm::dbgs()
         << type
         << " unhandled: can only convert scalar or vector element type\n");
     return nullptr;
   }
   if (!arrayElemType)

mlir/test/Conversion/FuncToSPIRV/types-to-spirv.mlir

 
 } // end module
 // -----
 
 //===----------------------------------------------------------------------===//
 // Index type
 //===----------------------------------------------------------------------===//
 
-// The index type is always converted into i32.
+// The index type is always converted into i32 or i64, with i32 being the
+// default.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
 } {
 
 // CHECK-LABEL: spirv.func @index_type
 // CHECK-SAME: %{{.*}}: i32
 func.func @index_type(%arg0: index) { return }
 

 func.func @float_vector(
   %arg0: vector<2xf16>
 ) { return }
 
 // CHECK-LABEL: spirv.func @one_element_vector
 // CHECK-SAME: %{{.+}}: i32
 func.func @one_element_vector(%arg0: vector<1xi8>) { return }
 
+// CHECK-LABEL: spirv.func @index_vector
+// CHECK-SAME: %{{.*}}: vector<4xi32>
+func.func @index_vector(%arg0: vector<4xindex>) { return }
+
 } // end module
 
 // -----
 
 // Check that capabilities for scalar types affects vector types too: having
 // special capabilities means keep vector types untouched.
 module attributes {
   spirv.target_env = #spirv.target_env<

 // NOEMU-LABEL: func @memref_1bit_type
 // NOEMU-SAME: memref<4x8xi1, #spirv.storage_class<StorageBuffer>>
 // NOEMU-SAME: memref<4x8xi1, #spirv.storage_class<Function>>
 func.func @memref_1bit_type(
     %arg0: memref<4x8xi1, #spirv.storage_class<StorageBuffer>>,
     %arg1: memref<4x8xi1, #spirv.storage_class<Function>>
 ) { return }
 
+// CHECK-LABEL: func @memref_index_type
+// CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<4 x i32, stride=4> [0])>, StorageBuffer>
+// CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<4 x i32>)>, Function>
+func.func @memref_index_type(
+    %arg0: memref<4xindex, #spirv.storage_class<StorageBuffer>>,
+    %arg1: memref<4xindex, #spirv.storage_class<Function>>
+) { return }
+
 } // end module
 
 // -----
 
 // Reject memory spaces.
 module attributes {
   spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>

 // CHECK-LABEL: spirv.func @float_tensor_types
 // CHECK-SAME: !spirv.array<32 x f32>
 // CHECK-SAME: !spirv.array<32 x f32>
 func.func @float_tensor_types(
   %arg1: tensor<8x4xf32>,
   %arg2: tensor<8x4xf16>
 ) { return }
 
+
+// CHECK-LABEL: spirv.func @index_tensor_type
+// CHECK-SAME: %{{.*}}: !spirv.array<20 x i32>
+func.func @index_tensor_type(%arg0: tensor<4x5xindex>) { return }
+
 } // end module
 
 // -----
 
 // Check that dynamic shapes are not supported.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
 } {