[mlir][Vector] Add support for 0-D 'vector.shape_cast' lowering

[dcaballe]

Summary

This PR adds support for shape casting from and to 0-D vectors.

mlir/lib/Dialect/Vector/Transforms/LowerVectorShapeCast.cpp

LogicalResult matchAndRewrite(vector::ShapeCastOp op,
     SmallVector<int64_t> resIdx(resRank);
     Value result = rewriter.create<arith::ConstantOp>(
         loc, resultVectorType, rewriter.getZeroAttr(resultVectorType));
     for (int64_t i = 0; i < numElts; i++) {
       if (i != 0) {
         incIdx(srcIdx, sourceVectorType, srcRank - 1);
         incIdx(resIdx, resultVectorType, resRank - 1);
       }
-      Value e = rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx);
-      result = rewriter.create<vector::InsertOp>(loc, e, result, resIdx);
+
+      Value extract;
+      if (srcRank == 0) {
+        // 0-D vector special case
+        assert(srcIdx.empty() && "Unexpected indices for 0-D vector");
+        extract = rewriter.create<vector::ExtractElementOp>(
+            loc, op.getSourceVectorType().getElementType(), op.getSource());
+      } else {
+        extract =
+            rewriter.create<vector::ExtractOp>(loc, op.getSource(), srcIdx);
+      }
+
+      if (resRank == 0) {
+        // 0-D vector special case
+        assert(resIdx.empty() && "Unexpected indices for 0-D vector");
+        result = rewriter.create<vector::InsertElementOp>(loc, extract, result);
+      } else {
+        result =
+            rewriter.create<vector::InsertOp>(loc, extract, result, resIdx);
+      }
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 
 private:
   static void incIdx(SmallVector<int64_t> &idx, VectorType tp, int64_t r) {
     assert(0 <= r && r < tp.getRank());

mlir/test/Dialect/Vector/vector-shape-cast-lowering-transforms.mlir

-
 // RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @nop_shape_cast
 // CHECK-SAME: %[[A:.*]]: vector<16xf32>
 // CHECK:      return %[[A]] : vector<16xf32>
 func.func @nop_shape_cast(%arg0: vector<16xf32>) -> vector<16xf32> {
   %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<16xf32>
   return %0 : vector<16xf32>

 // CHECK: %[[T11:.*]] = vector.insert %[[T10]], %[[T9]] [1, 0, 2] : f32 into vector<2x1x3xf32>
 // CHECK: return %[[T11]] : vector<2x1x3xf32>
 
 func.func @shape_cast_1d3d(%arg0 : vector<6xf32>) -> vector<2x1x3xf32> {
   %s = vector.shape_cast %arg0 : vector<6xf32> to vector<2x1x3xf32>
   return %s : vector<2x1x3xf32>
 }
 
+// CHECK-LABEL:   func.func @shape_cast_0d1d(
+// CHECK-SAME:                               %[[VAL_0:.*]]: vector<f32>) -> vector<1xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant dense<0.000000e+00> : vector<1xf32>
+// CHECK:           %[[VAL_2:.*]] = vector.extractelement %[[VAL_0]][] : vector<f32>
+// CHECK:           %[[VAL_3:.*]] = vector.insert %[[VAL_2]], %[[VAL_1]] [0] : f32 into vector<1xf32>
+// CHECK:           return %[[VAL_3]] : vector<1xf32>
+// CHECK:         }
+
+func.func @shape_cast_0d1d(%arg0 : vector<f32>) -> vector<1xf32> {
+  %s = vector.shape_cast %arg0 : vector<f32> to vector<1xf32>
+  return %s : vector<1xf32>
+}
+
+// CHECK-LABEL:   func.func @shape_cast_1d0d(
+// CHECK-SAME:                               %[[VAL_0:.*]]: vector<1xf32>) -> vector<f32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant dense<0.000000e+00> : vector<f32>
+// CHECK:           %[[VAL_2:.*]] = vector.extract %[[VAL_0]][0] : vector<1xf32>
+// CHECK:           %[[VAL_3:.*]] = vector.insertelement %[[VAL_2]], %[[VAL_1]][] : vector<f32>
+// CHECK:           return %[[VAL_3]] : vector<f32>
+// CHECK:         }
+
+func.func @shape_cast_1d0d(%arg0 : vector<1xf32>) -> vector<f32> {
+  %s = vector.shape_cast %arg0 : vector<1xf32> to vector<f32>
+  return %s : vector<f32>
+}
+
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !transform.any_op):
-  %f = transform.structured.match ops{["func.func"]} in %module_op 
+  %f = transform.structured.match ops{["func.func"]} in %module_op
     : (!transform.any_op) -> !transform.any_op
 
   %f2 = transform.vector.lower_shape_cast %f
       : (!transform.any_op) -> !transform.any_op
 }