[MLIR][Vector] Remove 0-d corner case condition.

Author: harrisonGPU

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

@@ -97,9 +97,6 @@ struct TransferReadPermutationLowering
   matchAndRewriteMaskableOp(vector::TransferReadOp op,
                             MaskingOpInterface maskOp,
                             PatternRewriter &rewriter) const override {
-    // TODO: support 0-d corner case.
-    if (op.getTransferRank() == 0)
-      return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
     // TODO: Support transfer_read inside MaskOp case.
     if (maskOp)
       return rewriter.notifyMatchFailure(op, "Masked case not supported");
@@ -326,9 +323,6 @@ struct TransferOpReduceRank
   matchAndRewriteMaskableOp(vector::TransferReadOp op,
                             MaskingOpInterface maskOp,
                             PatternRewriter &rewriter) const override {
-    // TODO: support 0-d corner case.
-    if (op.getTransferRank() == 0)
-      return rewriter.notifyMatchFailure(op, "0-d corner case not supported");
     // TODO: support masked case.
     if (maskOp)
       return rewriter.notifyMatchFailure(op, "Masked case not supported");
@@ -518,7 +512,7 @@ struct VectorLoadToMemrefLoadLowering
   }
 };
 
-/// Replace a vector.store with a vector.extractelement + memref.store.
+/// Replace a 0-d vector.store with a vector.extractelement + memref.store.
 struct VectorStoreToMemrefStoreLowering
     : public OpRewritePattern<vector::StoreOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -530,9 +524,15 @@ struct VectorStoreToMemrefStoreLowering
       return rewriter.notifyMatchFailure(storeOp, "not single element vector");
 
     Value extracted;
-    SmallVector<int64_t> indices(vecType.getRank(), 0);
-    extracted = rewriter.create<vector::ExtractOp>(
-        storeOp.getLoc(), storeOp.getValueToStore(), indices);
+    if (vecType.getRank() == 0) {
+      // TODO: Unifiy once ExtractOp supports 0-d vectors.
+      extracted = rewriter.create<vector::ExtractElementOp>(
+          storeOp.getLoc(), storeOp.getValueToStore());
+    } else {
+      SmallVector<int64_t> indices(vecType.getRank(), 0);
+      extracted = rewriter.create<vector::ExtractOp>(
+          storeOp.getLoc(), storeOp.getValueToStore(), indices);
+    }
 
     rewriter.replaceOpWithNewOp<memref::StoreOp>(
         storeOp, extracted, storeOp.getBase(), storeOp.getIndices());

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

@@ -2971,8 +2971,9 @@ func.func @vector_store_op_0d(%memref : memref<200x100xf32>, %i : index, %j : in
 // CHECK-LABEL: func @vector_store_op_0d
 // CHECK: %[[val:.*]] = arith.constant dense<1.100000e+01> : vector<f32>
 // CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[val]] : vector<f32> to vector<1xf32>
-// CHECK: %[[cast2:.*]] = builtin.unrealized_conversion_cast %[[cast]] : vector<1xf32> to f32
-// CHECK: memref.store %[[cast2]], %{{.*}}[%{{.*}}, %{{.*}}]
+// CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: %[[extracted:.*]] = llvm.extractelement %[[cast]][%[[c0]] : i64] : vector<1xf32>
+// CHECK: memref.store %[[extracted]], %{{.*}}[%{{.*}}, %{{.*}}]
 
 // -----
 

mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir

@@ -7,13 +7,15 @@ func.func @vector_transfer_ops_0d_memref(%mem: memref<f32>, %vec: vector<1x1x1xf
     %f0 = arith.constant 0.0 : f32
 
 //  CHECK-NEXT:   %[[S:.*]] = memref.load %[[MEM]][] : memref<f32>
+//  CHECK-NEXT:   %[[V:.*]] = vector.broadcast %[[S]] : f32 to vector<f32>
     %0 = vector.transfer_read %mem[], %f0 : memref<f32>, vector<f32>
 
-//  CHECK-NEXT:   memref.store %[[S]], %[[MEM]][] : memref<f32>
+//  CHECK-NEXT:   %[[SS:.*]] = vector.extractelement %[[V]][] : vector<f32>
+//  CHECK-NEXT:   memref.store %[[SS]], %[[MEM]][] : memref<f32>
     vector.transfer_write %0, %mem[] : vector<f32>, memref<f32>
 
-//  CHECK-NEXT:   %[[V:.*]] = vector.extract %arg1[0, 0, 0] : f32 from vector<1x1x1xf32>
-//  CHECK-NEXT:   memref.store %[[V]], %[[MEM]][] : memref<f32>
+//  CHECK-NEXT:   %[[VV:.*]] = vector.extract %arg1[0, 0, 0] : f32 from vector<1x1x1xf32>
+//  CHECK-NEXT:   memref.store %[[VV]], %[[MEM]][] : memref<f32>
     vector.store %vec, %mem[] : memref<f32>, vector<1x1x1xf32>
 
     return