[mlir] [VectorOps] Replace zero-scalar + splat into direct zero vector constant

[aartbik]

Summary

The scalar zero + splat yields more intermediate code than the direct
dense zero constant, and ultimately is lowered to exactly the same
LLVM IR operations, so no point wasting the intermediate code.

[rriddle]

There were no test cases changed by this?

[aartbik]

In D79758#2058101, @rriddle wrote:

There were no test cases changed by this?

No. I was a bit surprised too I did not even have to change even one, but all tests looked at a later lowered stage, which is identical.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::ExtractSlicesOp op,
     return success();
   }
 };
 
 /// Progressive lowering of InsertSlicesOp to series of InsertStridedSliceOp.
 /// One:
 ///   %x = vector.insert_slices %0
 /// is replaced by:
-///   %r0 = vector.splat 0
-//    %t1 = vector.tuple_get %0, 0
+///   %r0 = zero-result
+///   %t1 = vector.tuple_get %0, 0
 ///   %r1 = vector.insert_strided_slice %r0, %t1
-//    %t2 = vector.tuple_get %0, 1
+///   %t2 = vector.tuple_get %0, 1
 ///   %r2 = vector.insert_strided_slice %r1, %t2
 ///   ..
 ///   %x  = ..
 class InsertSlicesOpLowering : public OpRewritePattern<vector::InsertSlicesOp> {
 public:
   using OpRewritePattern<vector::InsertSlicesOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::InsertSlicesOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType vectorType = op.getResultVectorType();
     auto shape = vectorType.getShape();
 
     SmallVector<int64_t, 4> sizes;
     op.getSizes(sizes);
     SmallVector<int64_t, 4> strides;
     op.getStrides(strides); // all-ones at the moment
 
     // Prepare result.
-    auto elemType = vectorType.getElementType();
-    Value zero = rewriter.create<ConstantOp>(loc, elemType,
-                                             rewriter.getZeroAttr(elemType));
-    Value result = rewriter.create<SplatOp>(loc, vectorType, zero);
+    Value result = rewriter.create<ConstantOp>(
+        loc, vectorType, rewriter.getZeroAttr(vectorType));
 
     // For each element in the tuple, extract the proper strided slice.
     TupleType tupleType = op.getSourceTupleType();
     int64_t tupleSize = tupleType.size();
     auto sliceStrides = computeStrides(shape, sizes);
     for (int64_t i = 0; i < tupleSize; ++i) {
       auto vectorOffsets = delinearize(sliceStrides, i);
       auto elementOffsets =
if (srcRank < dstRank) {
         rewriter.replaceOpWithNewOp<SplatOp>(op, dstType, op.source());
         return success();
       }
       // Duplication.
       VectorType resType =
           VectorType::get(dstType.getShape().drop_front(), eltType);
       Value bcst =
           rewriter.create<vector::BroadcastOp>(loc, resType, op.source());
-      Value zero = rewriter.create<ConstantOp>(loc, eltType,
-                                               rewriter.getZeroAttr(eltType));
-      Value result = rewriter.create<SplatOp>(loc, dstType, zero);
+      Value result = rewriter.create<ConstantOp>(loc, dstType,
+                                                 rewriter.getZeroAttr(dstType));
       for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
         result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
       rewriter.replaceOp(op, result);
       return success();
     }
 
     // Find non-matching dimension, if any.
     assert(srcRank == dstRank);
LogicalResult matchAndRewrite(vector::BroadcastOp op,
     // becomes:
     //   %u = broadcast %y[0][0] : vector<2xf32> to vector <2x2xf32>
     //   %v = broadcast %y[1][0] : vector<2xf32> to vector <2x2xf32>
     //   %a = [%u, %v]
     //   ..
     //   %x = [%a,%b,%c,%d]
     VectorType resType =
         VectorType::get(dstType.getShape().drop_front(), eltType);
-    Value zero = rewriter.create<ConstantOp>(loc, eltType,
-                                             rewriter.getZeroAttr(eltType));
-    Value result = rewriter.create<SplatOp>(loc, dstType, zero);
+    Value result = rewriter.create<ConstantOp>(loc, dstType,
+                                               rewriter.getZeroAttr(dstType));
     if (m == 0) {
       // Stetch at start.
       Value ext = rewriter.create<vector::ExtractOp>(loc, op.source(), 0);
       Value bcst = rewriter.create<vector::BroadcastOp>(loc, resType, ext);
       for (int64_t d = 0, dim = dstType.getDimSize(0); d < dim; ++d)
         result = rewriter.create<vector::InsertOp>(loc, bcst, result, d);
     } else {
       // Stetch not at start.

 public:
   using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType resType = op.getResultType();
-    Type eltType = resType.getElementType();
 
     // Set up convenience transposition table.
     SmallVector<int64_t, 4> transp;
     for (auto attr : op.transp())
       transp.push_back(attr.cast<IntegerAttr>().getInt());
 
     // Generate fully unrolled extract/insert ops.
-    Value zero = rewriter.create<ConstantOp>(loc, eltType,
-                                             rewriter.getZeroAttr(eltType));
-    Value result = rewriter.create<SplatOp>(loc, resType, zero);
+    Value result = rewriter.create<ConstantOp>(loc, resType,
+                                               rewriter.getZeroAttr(resType));
     SmallVector<int64_t, 4> lhs(transp.size(), 0);
     SmallVector<int64_t, 4> rhs(transp.size(), 0);
     rewriter.replaceOp(op, expandIndices(loc, resType, 0, transp, lhs, rhs,
                                          op.vector(), result, rewriter));
     return success();
   }
 
 private:
LogicalResult matchAndRewrite(vector::OuterProductOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType rhsType = op.getOperandVectorTypeRHS();
     VectorType resType = op.getVectorType();
     Type eltType = resType.getElementType();
     Value acc = (op.acc().empty()) ? nullptr : op.acc()[0];
 
-    Value zero = rewriter.create<ConstantOp>(loc, eltType,
-                                             rewriter.getZeroAttr(eltType));
-    Value result = rewriter.create<SplatOp>(loc, resType, zero);
+    Value result = rewriter.create<ConstantOp>(loc, resType,
+                                               rewriter.getZeroAttr(resType));
     for (int64_t d = 0, e = resType.getDimSize(0); d < e; ++d) {
       auto pos = rewriter.getI64ArrayAttr(d);
       Value x = rewriter.create<vector::ExtractOp>(loc, eltType, op.lhs(), pos);
       Value b = rewriter.create<vector::BroadcastOp>(loc, rhsType, x);
       Value m;
       if (acc) {
         Value z = rewriter.create<vector::ExtractOp>(loc, rhsType, acc, pos);
         m = rewriter.create<vector::FMAOp>(loc, b, op.rhs(), z);
Value lowerParallel(vector::ContractionOp op, int64_t lhsIndex,
     lowIndexingMaps.push_back(adjustMap(iMap[0], iterIndex, rewriter));
     lowIndexingMaps.push_back(adjustMap(iMap[1], iterIndex, rewriter));
     lowIndexingMaps.push_back(adjustMap(iMap[2], iterIndex, rewriter));
     auto lowAffine = rewriter.getAffineMapArrayAttr(lowIndexingMaps);
     auto lowIter =
         rewriter.getArrayAttr(adjustIter(op.iterator_types(), iterIndex));
     // Unroll into a series of lower dimensional vector.contract ops.
     Location loc = op.getLoc();
-    Value result = zeroVector(loc, resType, rewriter);
+    Value result = rewriter.create<ConstantOp>(loc, resType,
+                                               rewriter.getZeroAttr(resType));
     for (int64_t d = 0; d < dimSize; ++d) {
       auto lhs = reshapeLoad(loc, op.lhs(), lhsType, lhsIndex, d, rewriter);
       auto rhs = reshapeLoad(loc, op.rhs(), rhsType, rhsIndex, d, rewriter);
       auto acc = reshapeLoad(loc, op.acc(), resType, resIndex, d, rewriter);
       Value lowContract = rewriter.create<vector::ContractionOp>(
           loc, lhs, rhs, acc, lowAffine, lowIter);
       result = reshapeStore(loc, lowContract, result, resType, resIndex, d,
                             rewriter);
Value lowerReduction(vector::ContractionOp op,
     assert(lookupRhs.hasValue() && "missing RHS parallel index");
     int64_t lhsIndex = lookupLhs.getValue();
     int64_t rhsIndex = lookupRhs.getValue();
     int64_t dimSize = lhsType.getDimSize(lhsIndex);
     assert(dimSize == rhsType.getDimSize(rhsIndex) && "corrupt shape");
     // Base case.
     if (lhsType.getRank() == 1) {
       assert(rhsType.getRank() == 1 && "corrupt contraction");
-      Value zero = zeroVector(loc, lhsType, rewriter);
+      Value zero = rewriter.create<ConstantOp>(loc, lhsType,
+                                               rewriter.getZeroAttr(lhsType));
       Value fma = rewriter.create<vector::FMAOp>(loc, op.lhs(), op.rhs(), zero);
       StringAttr kind = rewriter.getStringAttr("add");
       return rewriter.create<vector::ReductionOp>(loc, resType, kind, fma,
                                                   op.acc());
     }
     // Construct new iterator types and affine map array attribute.
     SmallVector<AffineMap, 4> lowIndexingMaps;
     lowIndexingMaps.push_back(adjustMap(iMap[0], iterIndex, rewriter));
for (int64_t d = 0; d < dimSize; ++d) {
       auto lhs = reshapeLoad(loc, op.lhs(), lhsType, lhsIndex, d, rewriter);
       auto rhs = reshapeLoad(loc, op.rhs(), rhsType, rhsIndex, d, rewriter);
       result = rewriter.create<vector::ContractionOp>(loc, lhs, rhs, result,
                                                       lowAffine, lowIter);
     }
     return result;
   }
 
-  // Helper method to construct a zero vector.
-  static Value zeroVector(Location loc, VectorType vType,
-                          PatternRewriter &rewriter) {
-    Type eltType = vType.getElementType();
-    Value zero = rewriter.create<ConstantOp>(loc, eltType,
-                                             rewriter.getZeroAttr(eltType));
-    return rewriter.create<SplatOp>(loc, vType, zero);
-  }
-
   // Helper to find an index in an affine map.
   static Optional<int64_t> getResultIndex(AffineMap map, int64_t index) {
     for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
       int64_t idx = map.getResult(i).cast<AffineDimExpr>().getPosition();
       if (idx == index)
         return i;
     }
     return None;
static Value reshapeLoad(Location loc, Value val, VectorType type,
     // At extraction dimension?
     if (index == 0) {
       auto posAttr = rewriter.getI64ArrayAttr(pos);
       return rewriter.create<vector::ExtractOp>(loc, lowType, val, posAttr);
     }
     // Unroll leading dimensions.
     VectorType vType = lowType.cast<VectorType>();
     VectorType resType = adjustType(type, index).cast<VectorType>();
-    Value result = zeroVector(loc, resType, rewriter);
+    Value result = rewriter.create<ConstantOp>(loc, resType,
+                                               rewriter.getZeroAttr(resType));
     for (int64_t d = 0, e = resType.getDimSize(0); d < e; d++) {
       auto posAttr = rewriter.getI64ArrayAttr(d);
       Value ext = rewriter.create<vector::ExtractOp>(loc, vType, val, posAttr);
       Value load = reshapeLoad(loc, ext, vType, index - 1, pos, rewriter);
       result = rewriter.create<vector::InsertOp>(loc, resType, load, result,
                                                  posAttr);
     }
     return result;
public:
   LogicalResult matchAndRewrite(vector::ShapeCastOp op,
                                 PatternRewriter &rewriter) const override {
     auto sourceVectorType = op.getSourceVectorType();
     auto resultVectorType = op.getResultVectorType();
     if (sourceVectorType.getRank() != 2 || resultVectorType.getRank() != 1)
       return failure();
 
     auto loc = op.getLoc();
-    auto elemType = sourceVectorType.getElementType();
-    Value zero = rewriter.create<ConstantOp>(loc, elemType,
-                                             rewriter.getZeroAttr(elemType));
-    Value desc = rewriter.create<SplatOp>(loc, resultVectorType, zero);
+    Value desc = rewriter.create<ConstantOp>(
+        loc, resultVectorType, rewriter.getZeroAttr(resultVectorType));
     unsigned mostMinorVectorSize = sourceVectorType.getShape()[1];
     for (int64_t i = 0, e = sourceVectorType.getShape().front(); i != e; ++i) {
       Value vec = rewriter.create<vector::ExtractOp>(loc, op.source(), i);
       desc = rewriter.create<vector::InsertStridedSliceOp>(
           loc, vec, desc,
           /*offsets=*/i * mostMinorVectorSize, /*strides=*/1);
     }
     rewriter.replaceOp(op, desc);
public:
   LogicalResult matchAndRewrite(vector::ShapeCastOp op,
                                 PatternRewriter &rewriter) const override {
     auto sourceVectorType = op.getSourceVectorType();
     auto resultVectorType = op.getResultVectorType();
     if (sourceVectorType.getRank() != 1 || resultVectorType.getRank() != 2)
       return failure();
 
     auto loc = op.getLoc();
-    auto elemType = sourceVectorType.getElementType();
-    Value zero = rewriter.create<ConstantOp>(loc, elemType,
-                                             rewriter.getZeroAttr(elemType));
-    Value desc = rewriter.create<SplatOp>(loc, resultVectorType, zero);
+    Value desc = rewriter.create<ConstantOp>(
+        loc, resultVectorType, rewriter.getZeroAttr(resultVectorType));
     unsigned mostMinorVectorSize = resultVectorType.getShape()[1];
     for (int64_t i = 0, e = resultVectorType.getShape().front(); i != e; ++i) {
       Value vec = rewriter.create<vector::ExtractStridedSliceOp>(
           loc, op.source(), /*offsets=*/i * mostMinorVectorSize,
           /*sizes=*/mostMinorVectorSize,
           /*strides=*/1);
       desc = rewriter.create<vector::InsertOp>(loc, vec, desc, i);
     }