[mlir][spirv] Support conversion of extract op from vector<1xT> type

[antiagainst]

mlir/test/Conversion/VectorToSPIRV/simple.mlir

38

Should we actually create another function to host this so that we can return the value to really CHECK against it?

[ThomasRaoux]

mlir/test/Conversion/VectorToSPIRV/simple.mlir

38

Moved it to a different function and check that the extract is replaced by the scalar source.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

matchAndRewrite(vector::ExtractOp extractOp, ArrayRef<Value> operands,
     if (resultVectorType && resultVectorType.getNumElements() > 1)
       return failure();
 
     auto dstType = getTypeConverter()->convertType(extractOp.getType());
     if (!dstType)
       return failure();
 
     vector::ExtractOp::Adaptor adaptor(operands);
+    if (adaptor.vector().getType().isa<spirv::ScalarType>()) {
+      rewriter.replaceOp(extractOp, adaptor.vector());
+      return success();
+    }
+
     int32_t id = getFirstIntValue(extractOp.position());
     rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
         extractOp, adaptor.vector(), id);
     return success();
   }
 };
 
 struct VectorExtractStridedSliceOpConvert final

mlir/test/Conversion/VectorToSPIRV/simple.mlir

 
 // CHECK-LABEL: func @extract
 //  CHECK-SAME: %[[ARG:.+]]: vector<2xf32>
 //       CHECK:   %{{.+}} = spv.CompositeExtract %[[ARG]][0 : i32] : vector<2xf32>
 //       CHECK:   %{{.+}} = spv.CompositeExtract %[[ARG]][1 : i32] : vector<2xf32>
 func @extract(%arg0 : vector<2xf32>) {
   %0 = "vector.extract"(%arg0) {position = [0]} : (vector<2xf32>) -> vector<1xf32>
   %1 = "vector.extract"(%arg0) {position = [1]} : (vector<2xf32>) -> f32
   spv.Return
 }
 
 // -----
 
+module attributes { spv.target_env = #spv.target_env<#spv.vce<v1.0, [Float16], []>, {}> } {
+
+// CHECK-LABEL: func @extract_scalar
+//  CHECK-SAME: %[[ARG0:.+]]: vector<2xf16>
+//  CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
+//       CHECK:   %[[S:.+]] = spv.Bitcast %[[ARG0]] : vector<2xf16> to f32
+//       CHECK:   spv.CompositeInsert %[[S]], %[[ARG1]][0 : i32] : f32 into vector<4xf32>
+func @extract_scalar(%arg0 : vector<2xf16>, %arg1 : vector<4xf32>) {
+  %0 = vector.bitcast %arg0 : vector<2xf16> to vector<1xf32>
+  %1 = vector.extract %0[0] : vector<1xf32>
+  %2 = vector.insert %1, %arg1[0] : f32 into vector<4xf32>
+  spv.Return
+}
+
+} // end module
+
+// -----
+
 // CHECK-LABEL: extract_insert
 //  CHECK-SAME: %[[V:.*]]: vector<4xf32>
 //       CHECK:   %[[S:.*]] = spv.CompositeExtract %[[V]][1 : i32] : vector<4xf32>
 //       CHECK:   spv.CompositeInsert %[[S]], %[[V]][0 : i32] : f32 into vector<4xf32>
 func @extract_insert(%arg0 : vector<4xf32>) {
   %0 = vector.extract %arg0[1] : vector<4xf32>
   %1 = vector.insert %0, %arg0[0] : f32 into vector<4xf32>
   spv.Return