[mlir][spirv] Allow usage of vector size 8 and 16 with Vector16 capability

[abialas]

Summary

Per spec, vector sizes 8 and 16 are allowed when Vector16 capability is present. This change expand the limitation of vector sizes to accept these sizes.

[antiagainst]

Cool! Can you also add a test to test/Dialect/SPIRV/availability.mlir to check that Vector16 capability kicks in?

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

175

This variable isn't needed really? I think we can just let valid sizes pass through to the final check?

mlir/test/Dialect/SPIRV/Serialization/ocl-ops.mlir

18

This is size 16? :)

[abialas]

Thank You for the review. Added requested changes, except testing in availability.mlir. As this is test type required capability it is not being kicked there when checking particular ops, similar to Int8 or Float16. I added one more test for this in vce-deduction.mlir

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

 // for the definition of the following types and type categories.
 
 def SPV_Void : TypeAlias<NoneType, "void">;
 def SPV_Bool : TypeAlias<I1, "bool">;
 def SPV_Integer : AnyIntOfWidths<[8, 16, 32, 64]>;
 def SPV_Int32 : TypeAlias<I32, "Int32">;
 def SPV_Float : FloatOfWidths<[16, 32, 64]>;
 def SPV_Float16or32 : FloatOfWidths<[16, 32]>;
-def SPV_Vector : VectorOfLengthAndType<[2, 3, 4],
+def SPV_Vector : VectorOfLengthAndType<[2, 3, 4, 8, 16],
                                        [SPV_Bool, SPV_Integer, SPV_Float]>;
 // Component type check is done in the type parser for the following SPIR-V
 // dialect-specific types so we use "Any" here.
 def SPV_AnyPtr : DialectType<SPIRV_Dialect, SPV_IsPtrType,
                              "any SPIR-V pointer type">;
 def SPV_AnyArray : DialectType<SPIRV_Dialect, SPV_IsArrayType,
                                "any SPIR-V array type">;
 def SPV_AnyCooperativeMatrix : DialectType<SPIRV_Dialect,

 def SPV_SignlessOrUnsignedInt : SignlessOrUnsignedIntOfWidths<[8, 16, 32, 64]>;
 
 class SPV_CoopMatrixOfType<list<Type> allowedTypes> :
   ContainerType<AnyTypeOf<allowedTypes>, SPV_IsCooperativeMatrixType,
     "$_self.cast<::mlir::spirv::CooperativeMatrixNVType>().getElementType()",
     "Cooperative Matrix">;
 
 class SPV_ScalarOrVectorOf<Type type> :
-    AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4], [type]>]>;
+    AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4, 8, 16], [type]>]>;
 
 class SPV_ScalarOrVectorOrCoopMatrixOf<Type type> :
-    AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4], [type]>,
+    AnyTypeOf<[type, VectorOfLengthAndType<[2, 3, 4, 8, 16], [type]>,
                SPV_CoopMatrixOfType<[type]>]>;
 
 def SPV_ScalarOrVector : AnyTypeOf<[SPV_Scalar, SPV_Vector]>;
 def SPV_ScalarOrVectorOrPtr : AnyTypeOf<[SPV_ScalarOrVector, SPV_AnyPtr]>;
 
 class SPV_Vec4<Type type> : VectorOfLengthAndType<[4], [type]>;
 def SPV_IntVec4 : SPV_Vec4<SPV_Integer>;
 def SPV_IOrUIVec4 : SPV_Vec4<SPV_SignlessOrUnsignedInt>;

mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp

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

 void CompositeType::getCapabilities(
     SPIRVType::CapabilityArrayRefVector &capabilities,
     Optional<StorageClass> storage) {
   TypeSwitch<Type>(*this)
       .Case<ArrayType, CooperativeMatrixNVType, MatrixType, RuntimeArrayType,
             StructType>(
           [&](auto type) { type.getCapabilities(capabilities, storage); })
       .Case<VectorType>([&](VectorType type) {
+        auto vecSize = getNumElements();
+        if (vecSize == 8 || vecSize == 16) {
+          static const Capability caps[] = {Capability::Vector16};
+          ArrayRef<Capability> ref(caps, llvm::array_lengthof(caps));
+          capabilities.push_back(ref);
+        }
         return type.getElementType().cast<ScalarType>().getCapabilities(
             capabilities, storage);
       })
       .Default([](Type) { llvm_unreachable("invalid composite type"); });
 }
 
 Optional<int64_t> CompositeType::getSizeInBytes() {
   if (auto arrayType = dyn_cast<ArrayType>())

mlir/test/Dialect/SPIRV/Serialization/ocl-ops.mlir

 // RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
 
 spv.module Physical64 OpenCL requires #spv.vce<v1.0, [Kernel, Addresses], []> {
   spv.func @float_insts(%arg0 : f32) "None" {
     // CHECK: {{%.*}} = spv.OCL.exp {{%.*}} : f32
     %0 = spv.OCL.exp %arg0 : f32
     // CHECK: {{%.*}} = spv.OCL.fabs {{%.*}} : f32
     %1 = spv.OCL.fabs %arg0 : f32
     spv.Return
   }
 
   spv.func @integer_insts(%arg0 : i32) "None" {
     // CHECK: {{%.*}} = spv.OCL.s_abs {{%.*}} : i32
     %0 = spv.OCL.s_abs %arg0 : i32
     spv.Return
   }
+  
+  spv.func @vector_size16(%arg0 : vector<16xf32>) "None" {
+    // CHECK: {{%.*}} = spv.OCL.fabs {{%.*}} : vector<16xf32>
+    %0 = spv.OCL.fabs %arg0 : vector<16xf32>
+    spv.Return
+  }
 }

mlir/test/Dialect/SPIRV/Transforms/vce-deduction.mlir

spv.target_env = #spv.target_env<
     #spv.vce<v1.3, [Shader, Float16], []>, {}>
 } {
   spv.func @fadd_function(%val : f16) -> f16 "None" {
     %0 = spv.FAdd %val, %val : f16
     spv.ReturnValue %0: f16
   }
 }
 
+// Using 16-element vectors requires Vector16.
+// CHECK: requires #spv.vce<v1.0, [Vector16, Shader], []>
+spv.module Logical GLSL450 attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.3, [Shader, Vector16], []>, {}>
+} {
+  spv.func @iadd_v16_function(%val : vector<16xi32>) -> vector<16xi32> "None" {
+    %0 = spv.IAdd %val, %val : vector<16xi32>
+    spv.ReturnValue %0: vector<16xi32>
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // Extension
 //===----------------------------------------------------------------------===//
 
 // Test deducing minimal extensions.
 // spv.SubgroupBallotKHR requires the SPV_KHR_shader_ballot extension.
 
 // CHECK: requires #spv.vce<v1.0, [SubgroupBallotKHR, Shader], [SPV_KHR_shader_ballot]>

mlir/test/Dialect/SPIRV/ops.mlir

func @logicalUnary(%arg0 : i1)
   %0 = spv.LogicalNot : i1
   return
 }
 
 // -----
 
 func @logicalUnary(%arg0 : i32)
 {
-  // expected-error @+1 {{operand #0 must be bool or vector of bool values of length 2/3/4, but got 'i32'}}
+  // expected-error @+1 {{operand #0 must be bool or vector of bool values of length 2/3/4/8/16, but got 'i32'}}
   %0 = spv.LogicalNot %arg0 : i32
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // spv.MemoryBarrier