[mlir][spirv] Support conversion of `CopySignOp` to spirv for 1D vector with 1 element

[guraypp]

Summary

Conversion of CopySignOp to SPIRV is supported for scalar and vectors but not 1D vectors with 1 element (aka vector<1xf32>). This revisions adds supports this by treating them as scalars.

An alternative solution would be to allow 0D vectors for SPIRV, but the spec [0] strictly defines the vector type as non-0D.
"Vector: An ordered homogeneous collection of two or more scalars. Vector sizes are quite restrictive and dependent on the execution model."

[0] https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_types

[ThomasRaoux]

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

154

I think a simpler solution is to use the result of the type converter directly as it already knows to treat vector with 1 element as scalar. So I think you can just replace copySignOp.getType() here by type

[guraypp]

address comments

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

154

That's a simple fix. Thanks for pointing that out!

[antiagainst]

mlir/test/Conversion/MathToSPIRV/math-to-core-spirv.mlir

72

This is actually not 0-D vector? It's a 1-D 1-element vector?

[guraypp]

mlir/test/Conversion/MathToSPIRV/math-to-core-spirv.mlir

72

right, I altered change note accordingly. Thanks for pointing that!

mlir/lib/Conversion/MathToSPIRV/MathToSPIRV.cpp

matchAndRewrite(math::CopySignOp copySignOp, OpAdaptor adaptor,
     Type intType = rewriter.getIntegerType(bitwidth);
     uint64_t intValue = uint64_t(1) << (bitwidth - 1);
 
     Value signMask = rewriter.create<spirv::ConstantOp>(
         loc, intType, rewriter.getIntegerAttr(intType, intValue));
     Value valueMask = rewriter.create<spirv::ConstantOp>(
         loc, intType, rewriter.getIntegerAttr(intType, intValue - 1u));
 
     if (auto vectorType = copySignOp.getType().dyn_cast<VectorType>()) {
+      if (vectorType.getNumElements() != 1) {
         assert(vectorType.getRank() == 1);
         int count = vectorType.getNumElements();
         intType = VectorType::get(count, intType);
 
         SmallVector<Value> signSplat(count, signMask);
-      signMask =
-          rewriter.create<spirv::CompositeConstructOp>(loc, intType, signSplat);
+        signMask = rewriter.create<spirv::CompositeConstructOp>(loc, intType,
+                                                                signSplat);
 
         SmallVector<Value> valueSplat(count, valueMask);
         valueMask = rewriter.create<spirv::CompositeConstructOp>(loc, intType,
                                                                  valueSplat);
       }
+    }
 
     Value lhsCast =
         rewriter.create<spirv::BitcastOp>(loc, intType, adaptor.getLhs());
     Value rhsCast =
         rewriter.create<spirv::BitcastOp>(loc, intType, adaptor.getRhs());
 
     Value value = rewriter.create<spirv::BitwiseAndOp>(
         loc, intType, ValueRange{lhsCast, valueMask});

mlir/test/Conversion/MathToSPIRV/math-to-core-spirv.mlir

 func.func @copy_sign_tensor(%value: tensor<3x3xf32>, %sign: tensor<3x3xf32>) -> tensor<3x3xf32> {
   %0 = math.copysign %value, %sign : tensor<3x3xf32>
   return %0: tensor<3x3xf32>
 }
 
 // CHECK-LABEL: func @copy_sign_tensor
 // CHECK-NEXT:    math.copysign {{%.+}}, {{%.+}} : tensor<3x3xf32>
 // CHECK-NEXT:    return
+// -----
+
+module attributes { spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Float16, Int16], []>, #spirv.resource_limits<>> } {
+
+func.func @copy_sign_vector_0D(%value: vector<1xf16>, %sign: vector<1xf16>) -> vector<1xf16> {
+  %0 = math.copysign %value, %sign : vector<1xf16>
+  return %0: vector<1xf16>
+}
+
+}
+
+// CHECK-LABEL: func @copy_sign_vector_0D
+//  CHECK-SAME: (%[[VALUE:.+]]: vector<1xf16>, %[[SIGN:.+]]: vector<1xf16>)
+//       CHECK:   %[[CASTVAL:.+]] = builtin.unrealized_conversion_cast %[[VALUE]] : vector<1xf16> to f16
+//       CHECK:   %[[CASTSIGN:.+]] = builtin.unrealized_conversion_cast %[[SIGN]] : vector<1xf16> to f16
+//       CHECK:   %[[SMASK:.+]] = spirv.Constant -32768 : i16
+//       CHECK:   %[[VMASK:.+]] = spirv.Constant 32767 : i16
+//       CHECK:   %[[VCAST:.+]] = spirv.Bitcast %[[CASTVAL]] : f16 to i16
+//       CHECK:   %[[SCAST:.+]] = spirv.Bitcast %[[CASTSIGN]] : f16 to i16
+//       CHECK:   %[[VAND:.+]] = spirv.BitwiseAnd %[[VCAST]], %[[VMASK]] : i16
+//       CHECK:   %[[SAND:.+]] = spirv.BitwiseAnd %[[SCAST]], %[[SMASK]] : i16
+//       CHECK:   %[[OR:.+]] = spirv.BitwiseOr %[[VAND]], %[[SAND]] : i16
+//       CHECK:   %[[RESULT:.+]] = spirv.Bitcast %[[OR]] : i16 to f16
+//       CHECK:   %[[CASTRESULT:.+]] = builtin.unrealized_conversion_cast %[[RESULT]] : f16 to vector<1xf16>
+//       CHECK:   return %[[CASTRESULT]]