[mlir][spirv] Support sub-byte element vector extension/truncation

[antiagainst]

Summary

This would mean allowing vector type conversion involving sub-byte
element types.

[kuhar]

mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp

104–113

nit: consider replacing this with return llvm::is_contained({2, 3, 4, 8, 16}, type.getNumElements());

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

320–322

hasValidSize permits 8 and 16 elements, so I don't understand the debug message here

[antiagainst]

Address comments

[kuhar]

LGTM

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTypes.h

 // SPIR-V composite type: VectorType, SPIR-V ArrayType, or SPIR-V StructType.
 class CompositeType : public SPIRVType {
 public:
   using SPIRVType::SPIRVType;
 
   static bool classof(Type type);
 
   /// Returns true if the given vector type is valid for the SPIR-V dialect.
+  static bool hasValidSize(VectorType);
   static bool isValid(VectorType);
 
   /// Return the number of elements of the type. This should only be called if
   /// hasCompileTimeKnownNumElements is true.
   unsigned getNumElements() const;
 
   Type getElementType(unsigned) const;
 

mlir/lib/Dialect/SPIRV/IR/SPIRVTypes.cpp

bool CompositeType::classof(Type type) {
   if (auto vectorType = llvm::dyn_cast<VectorType>(type))
     return isValid(vectorType);
   return llvm::isa<spirv::ArrayType, spirv::CooperativeMatrixType,
                    spirv::CooperativeMatrixNVType, spirv::JointMatrixINTELType,
                    spirv::MatrixType, spirv::RuntimeArrayType,
                    spirv::StructType>(type);
 }
 
-bool CompositeType::isValid(VectorType type) {
+bool CompositeType::hasValidSize(VectorType type) {
   switch (type.getNumElements()) {
   case 2:
   case 3:
   case 4:
   case 8:
   case 16:
-    break;
+    return true;
   default:
     return false;
   }
-  return type.getRank() == 1 && llvm::isa<ScalarType>(type.getElementType());
+}
+
+bool CompositeType::isValid(VectorType type) {
+  return type.getRank() == 1 && hasValidSize(type) &&
+         llvm::isa<ScalarType>(type.getElementType());
 }
 
 Type CompositeType::getElementType(unsigned index) const {
   return TypeSwitch<Type, Type>(*this)
       .Case<ArrayType, CooperativeMatrixType, CooperativeMatrixNVType,
             JointMatrixINTELType, RuntimeArrayType, VectorType>(
           [](auto type) { return type.getElementType(); })
       .Case<MatrixType>([](MatrixType type) { return type.getColumnType(); })

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

 /// 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 = {}) {
   type = cast<VectorType>(convertIndexElementType(type, options));
   auto scalarType = dyn_cast_or_null<spirv::ScalarType>(type.getElementType());
   if (!scalarType) {
+    // If this is not a spec allowed scalar type, try to handle sub-byte integer
+    // types.
+    auto intType = dyn_cast<IntegerType>(type.getElementType());
+    if (!intType) {
       LLVM_DEBUG(llvm::dbgs()
-               << type << " illegal: cannot convert non-scalar element type\n");
+                 << type
+                 << " illegal: cannot convert non-scalar element type\n");
       return nullptr;
     }
 
+    Type elementType = convertSubByteIntegerType(options, intType);
+    if (type.getRank() <= 1 && type.getNumElements() == 1)
+      return elementType;
+
+    if (!spirv::CompositeType::hasValidSize(type)) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << type << " illegal: > 4-element unimplemented\n");
+      return nullptr;
+    }
+
+    return VectorType::get(type.getShape(), elementType);
+  }
+
   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;
   }
 

mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir

 // CHECK-LABEL: @fpext2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fpext2(%arg0 : f32) -> f64 {
   // CHECK-NEXT: spirv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f32 to f64
   return %0: f64
 }
 
-// CHECK-LABEL: @trunci4
+// CHECK-LABEL: @trunci4_scalar
 //  CHECK-SAME: %[[ARG:.*]]: i32
-func.func @trunci4(%arg0 : i32) -> i4 {
+func.func @trunci4_scalar(%arg0 : i32) -> i4 {
   // CHECK: %[[MASK:.+]] = spirv.Constant 15 : i32
   // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[ARG]], %[[MASK]] : i32
   %0 = arith.trunci %arg0 : i32 to i4
   // CHECK: %[[RET:.+]] = builtin.unrealized_conversion_cast %[[AND]] : i32 to i4
   // CHECK: return %[[RET]] : i4
   return %0 : i4
 }
 
-// CHECK-LABEL: @zexti4
-func.func @zexti4(%arg0: i4) -> i32 {
+// CHECK-LABEL: @trunci4_vector
+//  CHECK-SAME: %[[ARG:.*]]: vector<2xi32>
+func.func @trunci4_vector(%arg0 : vector<2xi32>) -> vector<2xi4> {
+  // CHECK: %[[MASK:.+]] = spirv.Constant dense<15> : vector<2xi32>
+  // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[ARG]], %[[MASK]] : vector<2xi32>
+  %0 = arith.trunci %arg0 : vector<2xi32> to vector<2xi4>
+  // CHECK: %[[RET:.+]] = builtin.unrealized_conversion_cast %[[AND]] : vector<2xi32> to vector<2xi4>
+  // CHECK: return %[[RET]] : vector<2xi4>
+  return %0 : vector<2xi4>
+}
+
+// CHECK-LABEL: @zexti4_scalar
+func.func @zexti4_scalar(%arg0: i4) -> i32 {
   // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %{{.+}} : i4 to i32
   // CHECK: %[[MASK:.+]] = spirv.Constant 15 : i32
   // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[INPUT]], %[[MASK]] : i32
   %0 = arith.extui %arg0 : i4 to i32
   // CHECK: return %[[AND]] : i32
   return %0 : i32
 }
 
-// CHECK-LABEL: @sexti4
-func.func @sexti4(%arg0: i4) -> i32 {
+// CHECK-LABEL: @zexti4_vector
+func.func @zexti4_vector(%arg0: vector<3xi4>) -> vector<3xi32> {
+  // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %{{.+}} : vector<3xi4> to vector<3xi32>
+  // CHECK: %[[MASK:.+]] = spirv.Constant dense<15> : vector<3xi32>
+  // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[INPUT]], %[[MASK]] : vector<3xi32>
+  %0 = arith.extui %arg0 : vector<3xi4> to vector<3xi32>
+  // CHECK: return %[[AND]] : vector<3xi32>
+  return %0 : vector<3xi32>
+}
+
+// CHECK-LABEL: @sexti4_scalar
+func.func @sexti4_scalar(%arg0: i4) -> i32 {
   // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %arg0 : i4 to i32
   // CHECK: %[[SIZE:.+]] = spirv.Constant 28 : i32
   // CHECK: %[[SL:.+]] = spirv.ShiftLeftLogical %[[INPUT]], %[[SIZE]] : i32, i32
   // CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[SL]], %[[SIZE]] : i32, i32
   %0 = arith.extsi %arg0 : i4 to i32
   // CHECK: return %[[SR]] : i32
   return %0 : i32
 }
 
+// CHECK-LABEL: @sexti4_vector
+func.func @sexti4_vector(%arg0: vector<4xi4>) -> vector<4xi32> {
+  // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %arg0 : vector<4xi4> to vector<4xi32>
+  // CHECK: %[[SIZE:.+]] = spirv.Constant dense<28> : vector<4xi32>
+  // CHECK: %[[SL:.+]] = spirv.ShiftLeftLogical %[[INPUT]], %[[SIZE]] : vector<4xi32>, vector<4xi32>
+  // CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[SL]], %[[SIZE]] : vector<4xi32>, vector<4xi32>
+  %0 = arith.extsi %arg0 : vector<4xi4> to vector<4xi32>
+  // CHECK: return %[[SR]] : vector<4xi32>
+  return %0 : vector<4xi32>
+}
+
 } // end module
 
 // -----
 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Float16], []>, #spirv.resource_limits<>>

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

 } {
 
 // CHECK: spirv.func @integer2(%{{.+}}: i32)
 func.func @integer2(%arg0: i8) { return }
 
 // CHECK: spirv.func @integer4(%{{.+}}: i32)
 func.func @integer4(%arg0: i4) { return }
 
+// CHECK: spirv.func @v3i4(%{{.+}}: vector<3xi32>)
+func.func @v3i4(%arg0: vector<3xi4>) { return }
+
 } // end module
 
 // -----
 
 // Check that other bitwidths are not supported.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
 } {