[mlir][vector] Support 0-D vector when eliding single element reduction

[Lewuathe]

Summary

ElideSingleElementReduction causes assertion failure when we give 0-D vector. It's possible to fold the case by using vector.extractelement op instead. It's originally reported in https://github.com/llvm/llvm-project/issues/60193.

[Lewuathe]

Remove debug header.

[Lewuathe]

Add new line at the end of the test file.

[dcaballe]

Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

struct ElideSingleElementReduction : public OpRewritePattern<ReductionOp> {
   LogicalResult matchAndRewrite(ReductionOp reductionOp,
                                 PatternRewriter &rewriter) const override {
     // Masked reductions can't be folded until we can propagate the mask to the
     // resulting operation.
     auto maskableOp = cast<MaskableOpInterface>(reductionOp.getOperation());
     if (maskableOp.isMasked())
       return failure();
 
-    if (reductionOp.getVectorType().getDimSize(0) != 1)
+    auto vectorType = reductionOp.getVectorType();
+    if (vectorType.getRank() != 0 && vectorType.getDimSize(0) != 1)
       return failure();
 
     Location loc = reductionOp.getLoc();
-    Value result = rewriter.create<ExtractOp>(loc, reductionOp.getType(),
-                                              reductionOp.getVector(),
-                                              rewriter.getI64ArrayAttr(0));
+    Value result;
+    if (vectorType.getRank() == 0) {
+      result = rewriter.create<ExtractElementOp>(loc, reductionOp.getVector());
+    } else {
+      result = rewriter.create<ExtractOp>(loc, reductionOp.getType(),
+    }
 
     if (Value acc = reductionOp.getAcc())
       result = vector::makeArithReduction(rewriter, loc, reductionOp.getKind(),
                                           result, acc);
 
     rewriter.replaceOp(reductionOp, result);
     return success();
   }

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK-SAME:     %[[f:.*]]: f32
 //       CHECK:   return %[[f]]
 func.func @fold_extractelement_of_broadcast(%f: f32) -> f32 {
   %0 = vector.broadcast %f : f32 to vector<15xf32>
   %c5 = arith.constant 5 : index
   %1 = vector.extractelement %0 [%c5 : index] : vector<15xf32>
   return %1 : f32
 }
+
+// -----
+
+// CHECK-LABEL: func.func @fold_0d_vector_reduction
+func.func @fold_0d_vector_reduction(%arg0: vector<f32>) -> f32 {
+  // CHECK-NEXT: %[[RES:.*]] = vector.extractelement %arg{{.*}}[] : vector<f32>
+  // CHECK-NEXT: return %[[RES]] : f32
+  %0 = vector.reduction <add>, %arg0 : vector<f32> into f32
+  return %0 : f32
+}