[mlir][sparse] Canonicalizing arguments to genReshapeDstShape and foreachInSparseConstant

[wrengr]

Summary

These functions don't need a`PatternRewriter`, they only need an OpBuilder. And, the builder should be the first argument, before the Location, to match the style used everywhere else in MLIR.

[wrengr]

git-clang-format

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

 /// Generates the comparison `v != 0` where `v` is of numeric type.
 /// For floating types, we use the "unordered" comparator (i.e., returns
 /// true if `v` is NaN).
 Value genIsNonzero(OpBuilder &builder, Location loc, Value v);
 
 /// Computes the shape of destination tensor of a reshape operator. This is only
 /// used when operands have dynamic shape. The shape of the destination is
 /// stored into dstShape.
-void genReshapeDstShape(Location loc, PatternRewriter &rewriter,
+void genReshapeDstShape(OpBuilder &builder, Location loc,
                         SmallVectorImpl<Value> &dstShape,
                         ArrayRef<Value> srcShape,
                         ArrayRef<StaticSize> staticDstShape,
                         ArrayRef<ReassociationIndices> reassociation);
 
 /// Reshape coordinates during a reshaping operation.
 void reshapeCvs(OpBuilder &builder, Location loc,
                 ArrayRef<ReassociationIndices> reassociation,

 /// %c2 = arith.constant 28
 /// %v2 = complex.constant (3.0, 4.0)
 /// callback({%c2}, %v2)
 ///
 /// %c3 = arith.constant 31
 /// %v3 = complex.constant (5.0, 6.0)
 /// callback({%c3}, %v3)
 void foreachInSparseConstant(
-    Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
-    AffineMap order, function_ref<void(ArrayRef<Value>, Value)> callback);
+    OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,
+    function_ref<void(ArrayRef<Value>, Value)> callback);
 
 /// Loads `size`-many values from the memref, which must have rank-1 and
 /// size greater-or-equal to `size`.  If the optional `(offsetIdx,offsetVal)`
 /// arguments are provided, then the `offsetVal` will be added to the
 /// `offsetIdx`-th value after loading.
 SmallVector<Value> loadAll(OpBuilder &builder, Location loc, size_t size,
                            Value mem, size_t offsetIdx = 0,
                            Value offsetVal = Value());

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

if (tp.isIntOrIndex())
     return builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
                                          zero);
   if (tp.dyn_cast<ComplexType>())
     return builder.create<complex::NotEqualOp>(loc, v, zero);
   llvm_unreachable("Non-numeric type");
 }
 
 void mlir::sparse_tensor::genReshapeDstShape(
-    Location loc, PatternRewriter &rewriter, SmallVectorImpl<Value> &dstShape,
+    OpBuilder &builder, Location loc, SmallVectorImpl<Value> &dstShape,
     ArrayRef<Value> srcShape, ArrayRef<StaticSize> staticDstShape,
     ArrayRef<ReassociationIndices> reassociation) {
   // Collapse shape.
   if (reassociation.size() < srcShape.size()) {
     unsigned start = 0;
     for (const auto &map : llvm::enumerate(reassociation)) {
-      auto dstDim = constantIndex(rewriter, loc, 1);
+      auto dstDim = constantIndex(builder, loc, 1);
       for (unsigned i = start; i < start + map.value().size(); i++) {
-        dstDim = rewriter.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
+        dstDim = builder.create<arith::MulIOp>(loc, dstDim, srcShape[i]);
       }
       dstShape.push_back(dstDim);
       start = start + map.value().size();
     }
     assert(start == srcShape.size());
     return;
   }
 
for (unsigned j = start; j < start + map.size(); j++) {
         // by dividing srcDim by the product of the static dimensions.
         StaticSize product = 1;
         for (unsigned k = start; k < start + map.size(); k++) {
           if (staticDstShape[k] != ShapedType::kDynamic) {
             product *= staticDstShape[k];
           }
         }
         // Compute the dynamic dimension size.
-        Value productVal = constantIndex(rewriter, loc, product);
+        Value productVal = constantIndex(builder, loc, product);
         Value dynamicSize =
-            rewriter.create<arith::DivUIOp>(loc, srcDim, productVal);
+            builder.create<arith::DivUIOp>(loc, srcDim, productVal);
         dstShape.push_back(dynamicSize);
       } else {
         // The expanded dimension is statically known.
-        dstShape.push_back(constantIndex(rewriter, loc, staticDstShape[j]));
+        dstShape.push_back(constantIndex(builder, loc, staticDstShape[j]));
       }
     }
     start = start + map.size();
   }
   assert(start == staticDstShape.size());
 }
 
 void mlir::sparse_tensor::reshapeCvs(
for (; isa<scf::ForOp>(op->getParentOp()) ||
          isa<scf::ParallelOp>(op->getParentOp()) ||
          isa<scf::IfOp>(op->getParentOp());
        op = op->getParentOp())
     ;
   return op;
 }
 
 void sparse_tensor::foreachInSparseConstant(
-    Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
-    AffineMap order, function_ref<void(ArrayRef<Value>, Value)> callback) {
+    OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,
+    function_ref<void(ArrayRef<Value>, Value)> callback) {
   const Dimension dimRank = getSparseTensorType(attr).getDimRank();
   const auto coordinates = attr.getIndices().getValues<IntegerAttr>();
   const auto values = attr.getValues().getValues<Attribute>();
 
   // This is like the `Element<V>` class in the runtime library, but for
   // MLIR attributes.  In the future we may want to move this out into
   // a proper class definition to help improve code legibility (e.g.,
   // `first` -> `coords`, `second` -> `value`) as well as being able
void sparse_tensor::foreachInSparseConstant(
 
   SmallVector<Value> cvs;
   cvs.reserve(dimRank);
   for (size_t i = 0, nse = values.size(); i < nse; i++) {
     // Remap coordinates.
     cvs.clear();
     for (Dimension d = 0; d < dimRank; d++) {
       auto crd = elems[i].first[d].getInt();
-      cvs.push_back(rewriter.create<arith::ConstantIndexOp>(loc, crd));
+      cvs.push_back(builder.create<arith::ConstantIndexOp>(loc, crd));
     }
     // Remap value.
     Value val;
     if (attr.getElementType().isa<ComplexType>()) {
       auto valAttr = elems[i].second.cast<ArrayAttr>();
-      val = rewriter.create<complex::ConstantOp>(loc, attr.getElementType(),
-                                                 valAttr);
+      val = builder.create<complex::ConstantOp>(loc, attr.getElementType(),
     } else {
       auto valAttr = elems[i].second.cast<TypedAttr>();
-      val = rewriter.create<arith::ConstantOp>(loc, valAttr);
+      val = builder.create<arith::ConstantOp>(loc, valAttr);
     }
     assert(val);
     callback(cvs, val);
   }
 }
 
 SmallVector<Value> sparse_tensor::loadAll(OpBuilder &builder, Location loc,
                                           size_t size, Value mem,

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
   Value iter = params.genBuffers(srcTp.withoutOrdering(), srcDimSizes)
                    .genNewCall(Action::kToIterator, adaptor.getSrc());
   // Start a new COO for the destination tensor.
   SmallVector<Value> dstDimSizes;
   if (dstTp.hasStaticDimShape())
     // Static "shapes" are in fact "sizes".
     fillDimShape(rewriter, loc, dstTp, dstDimSizes);
   else
-    genReshapeDstShape(loc, rewriter, dstDimSizes, srcDimSizes,
+    genReshapeDstShape(rewriter, loc, dstDimSizes, srcDimSizes,
                        dstTp.getDimShape(), op.getReassociationIndices());
   const Value coo =
       params.genBuffers(dstTp, dstDimSizes).genNewCall(Action::kEmptyCOO);
   const Value dstPerm = params.getDim2LvlMap();
   // Construct a while loop over the iterator.
   const Type iTp = rewriter.getIndexType();
   const Value srcDimCoords = genAlloca(rewriter, loc, srcTp.getDimRank(), iTp);
   const Value dstDimCoords = genAlloca(rewriter, loc, dstTp.getDimRank(), iTp);

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 static LogicalResult genForeachOnSparseConstant(ForeachOp op,
                                                 RewriterBase &rewriter,
                                                 SparseElementsAttr attr) {
   auto loc = op.getLoc();
   SmallVector<Value> reduc = op.getInitArgs();
 
   // Foreach on constant.
   foreachInSparseConstant(
-      loc, rewriter, attr, op.getOrder().value_or(AffineMap()),
+      rewriter, loc, attr, op.getOrder().value_or(AffineMap()),
       [&reduc, &rewriter, op](ArrayRef<Value> cvs, Value v) mutable {
         SmallVector<Value> args;
         args.append(cvs.begin(), cvs.end());
         args.push_back(v);
         args.append(reduc);
         // Clones the foreach op to get a copy of the loop body.
         auto cloned = cast<ForeachOp>(rewriter.clone(*op.getOperation()));
         assert(args.size() == cloned.getBody()->getNumArguments());
LogicalResult matchAndRewrite(ReshapeOp op,
     sizesForTensor(rewriter, srcSizes, loc, srcTp, srcTensor);
     SmallVector<Value> dstSizes;
     SmallVector<Value> dstDynSizes;
     if (dstTp.hasStaticDimShape()) {
       for (Dimension d : dstTp.getDimShape())
         dstSizes.push_back(constantIndex(rewriter, loc, d));
     } else {
       ArrayRef<DynSize> dstShape = dstTp.getDimShape();
-      genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
+      genReshapeDstShape(rewriter, loc, dstSizes, srcSizes, dstShape,
                          op.getReassociationIndices());
       for (auto [idx, shape] : llvm::enumerate(dstShape)) {
         if (shape == ShapedType::kDynamic)
           dstDynSizes.push_back(dstSizes[idx]);
       }
     }
     Value nnz = rewriter.create<NumberOfEntriesOp>(loc, srcTensor);
     // Only need a unordered COO buffer if input and output are not sorted