[mlir][spirv] Do not introduce vector<1xT> in UnifyAliasedResource

[kuhar]

Summary

1-element vectors are not valid in SPIR-V and fail Bitcast op verification.

[kuhar]

Add tests and commnets

[kuhar]

Add missing newline

[kuhar]

Fix test

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

if ((srcElemType.isIntOrFloat() && dstElemType.isIntOrFloat()) ||
       // If both the source and destination are vector types, we need to make
       // sure the scalar type is the same for composite construction later.
       if (auto srcElemVecType = dyn_cast<VectorType>(srcElemType))
         if (auto dstElemVecType = dyn_cast<VectorType>(dstElemType)) {
           if (srcElemVecType.getElementType() !=
               dstElemVecType.getElementType()) {
             int64_t count =
                 dstNumBytes / (srcElemVecType.getElementTypeBitWidth() / 8);
-            auto castType =
-                VectorType::get({count}, srcElemVecType.getElementType());
-            for (auto &c : components)
+
+            // Make sure not to create 1-element vectors, which are illegal in
+            // SPIR-V.
+            Type castType = srcElemVecType.getElementType();
+            if (count > 1)
+              castType = VectorType::get({count}, castType);
+
+            for (Value &c : components)
               c = rewriter.create<spirv::BitcastOp>(loc, castType, c);
           }
         }
       Value vectorValue = rewriter.create<spirv::CompositeConstructOp>(
           loc, vectorType, components);
 
       if (!isa<VectorType>(srcElemType))
         vectorValue =

mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir

 // CHECK:   %[[LD1:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<4xf16>
 // CHECK:   %[[BC0:.+]] = spirv.Bitcast %[[LD0]] : vector<4xf16> to vector<2xf32>
 // CHECK:   %[[BC1:.+]] = spirv.Bitcast %[[LD1]] : vector<4xf16> to vector<2xf32>
 // CHECK:   %[[CC:.+]] = spirv.CompositeConstruct %[[BC0]], %[[BC1]] : (vector<2xf32>, vector<2xf32>) -> vector<4xf32>
 // CHECK:   spirv.ReturnValue %[[CC]]
 
 // -----
 
+spirv.module Logical GLSL450 {
+  spirv.GlobalVariable @var01_v2f16 bind(0, 1) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf16>, stride=4> [0])>, StorageBuffer>
+  spirv.GlobalVariable @var01_v2f32 bind(0, 1) {aliased} : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+
+  spirv.func @aliased(%index: i32) -> vector<3xf32> "None" {
+    %c0 = spirv.Constant 0 : i32
+    %v0 = spirv.Constant dense<0.0> : vector<3xf32>
+    %addr0 = spirv.mlir.addressof @var01_v2f16 : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf16>, stride=4> [0])>, StorageBuffer>
+    %ac0 = spirv.AccessChain %addr0[%c0, %index] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf16>, stride=4> [0])>, StorageBuffer>, i32, i32
+    %value0 = spirv.Load "StorageBuffer" %ac0 : vector<2xf16>
+
+    %addr1 = spirv.mlir.addressof @var01_v2f32 : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>
+    %ac1 = spirv.AccessChain %addr1[%c0, %index] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf32>, stride=8> [0])>, StorageBuffer>, i32, i32
+    %value1 = spirv.Load "StorageBuffer" %ac1 : vector<2xf32>
+
+    %val0_as_f32 = spirv.Bitcast %value0 : vector<2xf16> to f32
+
+    %res = spirv.CompositeConstruct %val0_as_f32, %value1 : (f32, vector<2xf32>) -> vector<3xf32>
+
+    spirv.ReturnValue %res : vector<3xf32>
+  }
+}
+
+// CHECK-LABEL: spirv.module
+
+// CHECK: spirv.GlobalVariable @var01_v2f16 bind(0, 1) : !spirv.ptr<!spirv.struct<(!spirv.rtarray<vector<2xf16>, stride=4> [0])>, StorageBuffer>
+// CHECK: spirv.func @aliased
+
+// CHECK:     %[[LD0:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<2xf16>
+// CHECK:     %[[LD1:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<2xf16>
+// CHECK:     %[[LD2:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<2xf16>
+
+// CHECK-DAG: %[[ELEM0:.+]] = spirv.Bitcast %[[LD0]] : vector<2xf16> to f32
+// CHECK-DAG: %[[ELEM1:.+]] = spirv.Bitcast %[[LD1]] : vector<2xf16> to f32
+// CHECK-DAG: %[[ELEM2:.+]] = spirv.Bitcast %[[LD2]] : vector<2xf16> to f32
+
+// CHECK:     %[[RES:.+]] = spirv.CompositeConstruct %[[ELEM0]], %{{.+}} : (f32, vector<2xf32>) -> vector<3xf32>
+// CHECK:     spirv.ReturnValue %[[RES]] : vector<3xf32>
+
+// -----
+
 // Make sure we do not crash on function arguments.
 
 spirv.module Logical GLSL450 {
   spirv.func @main(%arg0: !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>) "None" {
     %cst0_i32 = spirv.Constant 0 : i32
     %0 = spirv.AccessChain %arg0[%cst0_i32, %cst0_i32] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>, i32, i32
     spirv.Return
   }
 }
 
 // CHECK-LABEL: spirv.module
 // CHECK-LABEL: spirv.func @main
 // CHECK-SAME:  (%{{.+}}: !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>) "None"
 // CHECK:       spirv.Return