Address comments

Author: tyb0807

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

@@ -250,6 +250,10 @@ LogicalResult isValidMaskedInputVector(ArrayRef<int64_t> shape,
 /// create sub vectors.
 /// 5. Insert the sub vectors back into the final vector.
 /// 6. Replace the original op with the new result.
+///
+/// Expects the operation to be unrolled to have at most 1 result. When there's
+/// no result, expects the caller to pass in the `vectorTy` to be able to get
+/// the unroll factor.
 using UnrollVectorOpFn =
     function_ref<Value(PatternRewriter &, Location, VectorType, int64_t)>;
 

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

@@ -65,7 +65,7 @@ struct UnrollScatter : OpRewritePattern<vector::ScatterOp> {
     Value maskVec = op.getMask();
     Value valueVec = op.getValueToStore();
 
-    // Get the vector type from one of the vector operands
+    // Get the vector type from one of the vector operands.
     VectorType vectorTy = dyn_cast<VectorType>(indexVec.getType());
     if (!vectorTy)
       return failure();
@@ -85,8 +85,8 @@ struct UnrollScatter : OpRewritePattern<vector::ScatterOp> {
                                          indexSubVec, maskSubVec, valueSubVec,
                                          op.getAlignmentAttr());
 
-      // Return a dummy value since unrollVectorOp expects a Value
-      return rewriter.create<ub::PoisonOp>(loc, subTy);
+      // Return a dummy value since unrollVectorOp expects a Value.
+      return Value();
     };
 
     return unrollVectorOp(op, rewriter, unrollScatterFn, vectorTy);

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -433,15 +433,12 @@ vector::unrollVectorValue(TypedValue<VectorType> vector,
 LogicalResult vector::unrollVectorOp(Operation *op, PatternRewriter &rewriter,
                                      vector::UnrollVectorOpFn unrollFn,
                                      VectorType vectorTy) {
-  // If vector type is not provided, get it from the result
+  // If vector type is not provided, get it from the result.
   if (!vectorTy) {
-    if (op->getNumResults() != 1)
-      return rewriter.notifyMatchFailure(
-          op, "expected single result when vector type not provided");
-
+    assert(op->getNumResults() == 1 &&
+           "expected single result when vector type not provided");
     vectorTy = dyn_cast<VectorType>(op->getResult(0).getType());
-    if (!vectorTy)
-      return rewriter.notifyMatchFailure(op, "expected vector type");
+    assert(vectorTy && "expected result to have vector type");
   }
 
   if (vectorTy.getRank() < 2)
@@ -454,7 +451,7 @@ LogicalResult vector::unrollVectorOp(Operation *op, PatternRewriter &rewriter,
 
   Location loc = op->getLoc();
 
-  // Only create result value if the operation produces results
+  // Only create result value if the operation produces results.
   Value result;
   if (op->getNumResults() > 0) {
     result = ub::PoisonOp::create(rewriter, loc, vectorTy);
@@ -465,7 +462,7 @@ LogicalResult vector::unrollVectorOp(Operation *op, PatternRewriter &rewriter,
   for (int64_t i = 0, e = vectorTy.getShape().front(); i < e; ++i) {
     Value subVector = unrollFn(rewriter, loc, subTy, i);
 
-    // Only insert if we have a result to build
+    // Only insert if we have a result to build.
     if (op->getNumResults() > 0) {
       result = vector::InsertOp::create(rewriter, loc, subVector, result, i);
     }