[mlir][sparse][NFC] Use RewriterBase/OpBuilder when possible

[springerm]

Summary

Most functions do not need a PatternRewriter or ConversionPatternRewriter.

Depends On D125361

[springerm]

The dependent revisions are not needed. This is standalone.

[aartbik]

Changes look good, and it is of course good to require the minimum required type. But is there a major advantage of doing it this way? (follow up or otherwise)?

[springerm]

In D125466#3509782, @aartbik wrote:

Changes look good, and it is of course good to require the minimum required type. But is there a major advantage of doing it this way? (follow up or otherwise)?

The API of One-Shot Bufferize provides a RewriterBase &, so this is needed in a few places to compose with the sparse compiler. Since I was already at it, I went through the entire sparsecompiler codebase and updated all PatternRewriter/ConversionPatternRewriter for consistency.

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

#endif
   /// Returns index of the root expression.
   unsigned buildLattices(unsigned e, unsigned i);
 
   /// Builds a tensor expression from the given Linalg operation.
   /// Returns index of the root expression on success.
   Optional<unsigned> buildTensorExpFromLinalg(linalg::GenericOp op);
 
   /// Rebuilds SSA format from a tensor expression.
-  Value buildExp(PatternRewriter &rewriter, Location loc, unsigned e, Value v0,
+  Value buildExp(RewriterBase &rewriter, Location loc, unsigned e, Value v0,
                  Value v1);
 
 private:
   /// Private helpers.
   bool maybeZero(unsigned e) const;
   bool isInvariant(unsigned e) const;
   Type inferType(unsigned e, Value src);
 

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

 enum class EmitCInterface : bool { Off = false, On = true };
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
 /// Returns the equivalent of `void*` for opaque arguments to the
 /// execution engine.
-static Type getOpaquePointerType(PatternRewriter &rewriter) {
-  return LLVM::LLVMPointerType::get(rewriter.getI8Type());
+static Type getOpaquePointerType(OpBuilder &builder) {
+  return LLVM::LLVMPointerType::get(builder.getI8Type());
 }
 
 /// Returns a function reference (first hit also inserts into module). Sets
 /// the "_emit_c_interface" on the function declaration when requested,
 /// so that LLVM lowering generates a wrapper function that takes care
 /// of ABI complications with passing in and returning MemRefs to C functions.
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name,
                                  TypeRange resultType, ValueRange operands,
static func::CallOp createFuncCall(OpBuilder &builder, Operation *op,
                                    ValueRange operands,
                                    EmitCInterface emitCInterface) {
   auto fn = getFunc(op, name, resultType, operands, emitCInterface);
   return builder.create<func::CallOp>(op->getLoc(), resultType, fn, operands);
 }
 
 /// Replaces the `op` with  a `CallOp` to the function reference returned
 /// by `getFunc()`.
-static func::CallOp replaceOpWithFuncCall(PatternRewriter &rewriter,
-                                          Operation *op, StringRef name,
-                                          TypeRange resultType,
+static func::CallOp replaceOpWithFuncCall(RewriterBase &rewriter, Operation *op,
+                                          StringRef name, TypeRange resultType,
                                           ValueRange operands,
                                           EmitCInterface emitCInterface) {
   auto fn = getFunc(op, name, resultType, operands, emitCInterface);
   return rewriter.replaceOpWithNewOp<func::CallOp>(op, resultType, fn,
                                                    operands);
 }
 
 /// Generates dimension size call.
-static Value genDimSizeCall(ConversionPatternRewriter &rewriter, Operation *op,
+static Value genDimSizeCall(OpBuilder &builder, Operation *op,
                             SparseTensorEncodingAttr &enc, Value src,
                             int64_t idx) {
   // Permute the index according to an optional dimension ordering.
   if (AffineMap p = enc.getDimOrdering())
     idx = p.getPermutedPosition(idx);
   // Generate the call.
   StringRef name = "sparseDimSize";
-  SmallVector<Value, 2> params{src, constantIndex(rewriter, op->getLoc(), idx)};
-  Type iTp = rewriter.getIndexType();
-  return createFuncCall(rewriter, op, name, iTp, params, EmitCInterface::Off)
+  SmallVector<Value, 2> params{src, constantIndex(builder, op->getLoc(), idx)};
+  Type iTp = builder.getIndexType();
+  return createFuncCall(builder, op, name, iTp, params, EmitCInterface::Off)
       .getResult(0);
 }
 
 /// Generates a call into the "swiss army knife" method of the sparse runtime
 /// support library for materializing sparse tensors into the computation.
-static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
+static Value genNewCall(OpBuilder &builder, Operation *op,
                         ArrayRef<Value> params) {
   StringRef name = "newSparseTensor";
-  Type pTp = getOpaquePointerType(rewriter);
-  return createFuncCall(rewriter, op, name, pTp, params, EmitCInterface::On)
+  Type pTp = getOpaquePointerType(builder);
+  return createFuncCall(builder, op, name, pTp, params, EmitCInterface::On)
       .getResult(0);
 }
 
 /// Populates given sizes array from type.
-static void sizesFromType(ConversionPatternRewriter &rewriter,
-                          SmallVector<Value, 4> &sizes, Location loc,
-                          ShapedType stp) {
+static void sizesFromType(OpBuilder &builder, SmallVector<Value, 4> &sizes,
+                          Location loc, ShapedType stp) {
   auto shape = stp.getShape();
   for (unsigned i = 0, rank = stp.getRank(); i < rank; i++) {
     uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i];
-    sizes.push_back(constantIndex(rewriter, loc, s));
+    sizes.push_back(constantIndex(builder, loc, s));
   }
 }
 
 /// Populates given sizes array from source.
-static void sizesFromSrc(ConversionPatternRewriter &rewriter,
-                         SmallVector<Value, 4> &sizes, Location loc,
-                         Value src) {
+static void sizesFromSrc(OpBuilder &builder, SmallVector<Value, 4> &sizes,
+                         Location loc, Value src) {
   unsigned rank = src.getType().cast<ShapedType>().getRank();
   for (unsigned i = 0; i < rank; i++)
-    sizes.push_back(linalg::createOrFoldDimOp(rewriter, loc, src, i));
+    sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, i));
 }
 
 /// Populates given sizes array from type (for static sizes) and from
 /// an already converted into opague pointer source (for dynamic sizes).
-static void sizesFromPtr(ConversionPatternRewriter &rewriter,
-                         SmallVector<Value, 4> &sizes, Operation *op,
-                         SparseTensorEncodingAttr &enc, ShapedType stp,
-                         Value src) {
+static void sizesFromPtr(OpBuilder &builder, SmallVector<Value, 4> &sizes,
+                         Operation *op, SparseTensorEncodingAttr &enc,
+                         ShapedType stp, Value src) {
   Location loc = op->getLoc();
   auto shape = stp.getShape();
   for (unsigned i = 0, rank = stp.getRank(); i < rank; i++)
     if (shape[i] == ShapedType::kDynamicSize)
-      sizes.push_back(genDimSizeCall(rewriter, op, enc, src, i));
+      sizes.push_back(genDimSizeCall(builder, op, enc, src, i));
     else
-      sizes.push_back(constantIndex(rewriter, loc, shape[i]));
+      sizes.push_back(constantIndex(builder, loc, shape[i]));
 }
 
 /// Generates an uninitialized temporary buffer of the given size and
 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
-static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
-                       Value sz, Type tp) {
+static Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp) {
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
-  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
+  return builder.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
 }
 
 /// Generates an uninitialized buffer of the given size and type,
 /// but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`). Unlike temporary buffers on the stack,
 /// this buffer must be explicitly deallocated by client.
-static Value genAlloc(ConversionPatternRewriter &rewriter, Location loc,
-                      Value sz, Type tp) {
+static Value genAlloc(RewriterBase &rewriter, Location loc, Value sz, Type tp) {
   auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
   return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{sz});
 }
 
 /// Generates an uninitialized temporary buffer of the given size and
 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
-static Value genAlloca(ConversionPatternRewriter &rewriter, Location loc,
-                       unsigned sz, Type tp) {
-  return genAlloca(rewriter, loc, constantIndex(rewriter, loc, sz), tp);
+static Value genAlloca(OpBuilder &builder, Location loc, unsigned sz, Type tp) {
+  return genAlloca(builder, loc, constantIndex(builder, loc, sz), tp);
 }
 
 /// Generates an uninitialized temporary buffer with room for one value
 /// of the given type, and returns the `memref<$tp>`.
-static Value genAllocaScalar(ConversionPatternRewriter &rewriter, Location loc,
-                             Type tp) {
-  return rewriter.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
+static Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp) {
+  return builder.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
 }
 
 /// Generates a temporary buffer of the given type and given contents.
-static Value genBuffer(ConversionPatternRewriter &rewriter, Location loc,
-                       ValueRange values) {
+static Value genBuffer(OpBuilder &builder, Location loc, ValueRange values) {
   unsigned sz = values.size();
   assert(sz >= 1);
-  Value buffer = genAlloca(rewriter, loc, sz, values[0].getType());
+  Value buffer = genAlloca(builder, loc, sz, values[0].getType());
   for (unsigned i = 0; i < sz; i++) {
-    Value idx = constantIndex(rewriter, loc, i);
-    rewriter.create<memref::StoreOp>(loc, values[i], buffer, idx);
+    Value idx = constantIndex(builder, loc, i);
+    builder.create<memref::StoreOp>(loc, values[i], buffer, idx);
   }
   return buffer;
 }
 
 /// Populates parameters required to call the "swiss army knife" method of the
 /// sparse runtime support library for materializing sparse tensors into the
 /// computation.
-static void newParams(ConversionPatternRewriter &rewriter,
-                      SmallVector<Value, 8> &params, Operation *op,
-                      ShapedType stp, SparseTensorEncodingAttr &enc,
-                      Action action, ValueRange szs, Value ptr = Value()) {
+static void newParams(OpBuilder &builder, SmallVector<Value, 8> &params,
+                      Operation *op, ShapedType stp,
+                      SparseTensorEncodingAttr &enc, Action action,
+                      ValueRange szs, Value ptr = Value()) {
   Location loc = op->getLoc();
   ArrayRef<SparseTensorEncodingAttr::DimLevelType> dlt = enc.getDimLevelType();
   unsigned sz = dlt.size();
   // Sparsity annotations.
   SmallVector<Value, 4> attrs;
   for (unsigned i = 0; i < sz; i++)
-    attrs.push_back(constantDimLevelTypeEncoding(rewriter, loc, dlt[i]));
-  params.push_back(genBuffer(rewriter, loc, attrs));
+    attrs.push_back(constantDimLevelTypeEncoding(builder, loc, dlt[i]));
+  params.push_back(genBuffer(builder, loc, attrs));
   // Dimension sizes array of the enveloping tensor. Useful for either
   // verification of external data, or for construction of internal data.
-  params.push_back(genBuffer(rewriter, loc, szs));
+  params.push_back(genBuffer(builder, loc, szs));
   // Dimension order permutation array. This is the "identity" permutation by
   // default, or otherwise the "reverse" permutation of a given ordering, so
   // that indices can be mapped quickly to the right position.
   SmallVector<Value, 4> rev(sz);
   if (AffineMap p = enc.getDimOrdering()) {
     for (unsigned i = 0; i < sz; i++)
-      rev[p.getDimPosition(i)] = constantIndex(rewriter, loc, i);
+      rev[p.getDimPosition(i)] = constantIndex(builder, loc, i);
   } else {
     for (unsigned i = 0; i < sz; i++)
-      rev[i] = constantIndex(rewriter, loc, i);
+      rev[i] = constantIndex(builder, loc, i);
   }
-  params.push_back(genBuffer(rewriter, loc, rev));
+  params.push_back(genBuffer(builder, loc, rev));
   // Secondary and primary types encoding.
   Type elemTp = stp.getElementType();
-  params.push_back(constantPointerTypeEncoding(rewriter, loc, enc));
-  params.push_back(constantIndexTypeEncoding(rewriter, loc, enc));
-  params.push_back(constantPrimaryTypeEncoding(rewriter, loc, elemTp));
+  params.push_back(constantPointerTypeEncoding(builder, loc, enc));
+  params.push_back(constantIndexTypeEncoding(builder, loc, enc));
+  params.push_back(constantPrimaryTypeEncoding(builder, loc, elemTp));
   // User action.
-  params.push_back(constantAction(rewriter, loc, action));
+  params.push_back(constantAction(builder, loc, action));
   // Payload pointer.
   if (!ptr)
-    ptr = rewriter.create<LLVM::NullOp>(loc, getOpaquePointerType(rewriter));
+    ptr = builder.create<LLVM::NullOp>(loc, getOpaquePointerType(builder));
   params.push_back(ptr);
 }
 
 /// Generates the code to read the value from tensor[ivs], and conditionally
 /// stores the indices ivs to the memory in ind. The generated code looks like
 /// the following and the insertion point after this routine is inside the
 /// if-then branch behind the assignment to ind. This is to ensure that the
 /// addEltX call generated after is inside the if-then branch.
 ///    if (tensor[ivs]!=0) {
 ///      ind = ivs
-static Value genIndexAndValueForDense(ConversionPatternRewriter &rewriter,
-                                      Location loc, Value tensor, Value ind,
-                                      ValueRange ivs) {
-  Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
-  Value cond = genIsNonzero(rewriter, loc, val);
-  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, cond, /*else*/ false);
-  rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
+static Value genIndexAndValueForDense(OpBuilder &builder, Location loc,
+                                      Value tensor, Value ind, ValueRange ivs) {
+  Value val = builder.create<tensor::ExtractOp>(loc, tensor, ivs);
+  Value cond = genIsNonzero(builder, loc, val);
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, cond, /*else*/ false);
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   unsigned i = 0;
   for (auto iv : ivs) {
-    Value idx = constantIndex(rewriter, loc, i++);
-    rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
+    Value idx = constantIndex(builder, loc, i++);
+    builder.create<memref::StoreOp>(loc, iv, ind, idx);
   }
   return val;
 }
 
 /// Generates a call to release/delete a `SparseTensorCOO`.
 static void genDelCOOCall(OpBuilder &builder, Operation *op, Type elemTp,
                           Value coo) {
   SmallString<21> name{"delSparseTensorCOO", primaryTypeFunctionSuffix(elemTp)};
   TypeRange noTp;
   createFuncCall(builder, op, name, noTp, coo, EmitCInterface::Off);
 }
 
 /// Generates a call that adds one element to a coordinate scheme.
 /// In particular, this generates code like the following:
 ///   val = a[i1,..,ik];
 ///   if val != 0
 ///     t->add(val, [i1,..,ik], [p1,..,pk]);
-static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
-                          Type eltType, Value ptr, Value val, Value ind,
-                          Value perm) {
+static void genAddEltCall(OpBuilder &builder, Operation *op, Type eltType,
+                          Value ptr, Value val, Value ind, Value perm) {
   SmallString<9> name{"addElt", primaryTypeFunctionSuffix(eltType)};
   SmallVector<Value, 4> params{ptr, val, ind, perm};
-  Type pTp = getOpaquePointerType(rewriter);
-  createFuncCall(rewriter, op, name, pTp, params, EmitCInterface::On);
+  Type pTp = getOpaquePointerType(builder);
+  createFuncCall(builder, op, name, pTp, params, EmitCInterface::On);
 }
 
 /// Generates a call to `iter->getNext()`.  If there is a next element,
 /// then it is copied into the out-parameters `ind` and `elemPtr`,
 /// and the return value is true.  If there isn't a next element, then
 /// the memory for `iter` is freed and the return value is false.
-static Value genGetNextCall(ConversionPatternRewriter &rewriter, Operation *op,
-                            Value iter, Value ind, Value elemPtr) {
+static Value genGetNextCall(OpBuilder &builder, Operation *op, Value iter,
+                            Value ind, Value elemPtr) {
   Type elemTp = elemPtr.getType().cast<ShapedType>().getElementType();
   SmallString<10> name{"getNext", primaryTypeFunctionSuffix(elemTp)};
   SmallVector<Value, 3> params{iter, ind, elemPtr};
-  Type i1 = rewriter.getI1Type();
-  return createFuncCall(rewriter, op, name, i1, params, EmitCInterface::On)
+  Type i1 = builder.getI1Type();
+  return createFuncCall(builder, op, name, i1, params, EmitCInterface::On)
       .getResult(0);
 }
 
 /// If the tensor is a sparse constant, generates and returns the pair of
 /// the constants for the indices and the values.
 static Optional<std::pair<Value, Value>>
-genSplitSparseConstant(ConversionPatternRewriter &rewriter, Location loc,
-                       Value tensor) {
+genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor) {
   if (auto constOp = tensor.getDefiningOp<arith::ConstantOp>()) {
     if (auto attr = constOp.getValue().dyn_cast<SparseElementsAttr>()) {
       DenseElementsAttr indicesAttr = attr.getIndices();
-      Value indices = rewriter.create<arith::ConstantOp>(loc, indicesAttr);
+      Value indices = builder.create<arith::ConstantOp>(loc, indicesAttr);
       DenseElementsAttr valuesAttr = attr.getValues();
-      Value values = rewriter.create<arith::ConstantOp>(loc, valuesAttr);
+      Value values = builder.create<arith::ConstantOp>(loc, valuesAttr);
       return std::make_pair(indices, values);
     }
   }
   return {};
 }
 
 /// Generates the code to copy the index at indices[ivs] to ind, and return
 /// the value at value[ivs].
-static Value genIndexAndValueForSparse(ConversionPatternRewriter &rewriter,
-                                       Location loc, Value indices,
-                                       Value values, Value ind, ValueRange ivs,
-                                       unsigned rank) {
+static Value genIndexAndValueForSparse(OpBuilder &builder, Location loc,
+                                       Value indices, Value values, Value ind,
+                                       ValueRange ivs, unsigned rank) {
   for (unsigned i = 0; i < rank; i++) {
-    Value idx = constantIndex(rewriter, loc, i);
-    Value val = rewriter.create<tensor::ExtractOp>(loc, indices,
-                                                   ValueRange{ivs[0], idx});
-    val =
-        rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), val);
-    rewriter.create<memref::StoreOp>(loc, val, ind, idx);
+    Value idx = constantIndex(builder, loc, i);
+    Value val = builder.create<tensor::ExtractOp>(loc, indices,
+    val = builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), val);
+    builder.create<memref::StoreOp>(loc, val, ind, idx);
   }
-  return rewriter.create<tensor::ExtractOp>(loc, values, ivs[0]);
+  return builder.create<tensor::ExtractOp>(loc, values, ivs[0]);
 }
 
 /// Generates code to allocate a tensor of the given type, and zero
 /// initialize it.  If the tensor type has any dynamic sizes, then the
 /// `sizes` parameter should be as filled by sizesFromPtr(); that way
 /// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
-static Value allocDenseTensor(ConversionPatternRewriter &rewriter, Location loc,
+static Value allocDenseTensor(OpBuilder &builder, Location loc,
                               RankedTensorType tensorTp, ValueRange sizes) {
   Type elemTp = tensorTp.getElementType();
   auto shape = tensorTp.getShape();
   auto memTp = MemRefType::get(shape, elemTp);
   SmallVector<Value> dynamicSizes;
   for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {
     if (shape[i] == ShapedType::kDynamicSize)
       dynamicSizes.push_back(sizes[i]);
   }
-  Value mem = rewriter.create<memref::AllocOp>(loc, memTp, dynamicSizes);
-  Value zero = constantZero(rewriter, loc, elemTp);
-  rewriter.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
+  Value mem = builder.create<memref::AllocOp>(loc, memTp, dynamicSizes);
+  Value zero = constantZero(builder, loc, elemTp);
+  builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{mem});
   return mem;
 }
 
 /// Inserts the element returned by genGetNextCall(_, ind, elemPtr) into
 /// the tensor created by allocDenseTensor().  The `rank` is the rank
 /// of the `tensor` and the length of `ind`.
-static void insertScalarIntoDenseTensor(ConversionPatternRewriter &rewriter,
-                                        Location loc, Value elemPtr,
-                                        Value tensor, unsigned rank,
-                                        Value ind) {
+static void insertScalarIntoDenseTensor(OpBuilder &builder, Location loc,
+                                        Value elemPtr, Value tensor,
+                                        unsigned rank, Value ind) {
   SmallVector<Value, 4> ivs;
   ivs.reserve(rank);
   for (unsigned i = 0; i < rank; i++) {
-    Value idx = constantIndex(rewriter, loc, i);
-    ivs.push_back(rewriter.create<memref::LoadOp>(loc, ind, idx));
+    Value idx = constantIndex(builder, loc, i);
+    ivs.push_back(builder.create<memref::LoadOp>(loc, ind, idx));
   }
-  Value elemV = rewriter.create<memref::LoadOp>(loc, elemPtr);
-  rewriter.create<memref::StoreOp>(loc, elemV, tensor, ivs);
+  Value elemV = builder.create<memref::LoadOp>(loc, elemPtr);
+  builder.create<memref::StoreOp>(loc, elemV, tensor, ivs);
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {

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

default:
     llvm_unreachable("unexpected reduction operator");
   }
 }
 
 /// Generates an initial value for a vector reduction, following the scheme
 /// given in Chapter 5 of "The Software Vectorization Handbook", where the
 /// initial scalar value is correctly embedded in the vector reduction value,
 /// and a straightforward horizontal reduction will complete the operation.
-static Value genVectorReducInit(CodeGen &codegen, PatternRewriter &rewriter,
+static Value genVectorReducInit(CodeGen &codegen, OpBuilder &builder,
                                 Location loc, VectorType vtp) {
   Value r = codegen.redVal;
   switch (codegen.redKind) {
   case kNoReduc:
     break;
   case kSum:
   case kXor:
     // Initialize reduction vector to: | 0 | .. | 0 | r |
-    return rewriter.create<vector::InsertElementOp>(
-        loc, r, constantZero(rewriter, loc, vtp),
-        constantIndex(rewriter, loc, 0));
+    return builder.create<vector::InsertElementOp>(
+        loc, r, constantZero(builder, loc, vtp),
+        constantIndex(builder, loc, 0));
   case kProduct:
     // Initialize reduction vector to: | 1 | .. | 1 | r |
-    return rewriter.create<vector::InsertElementOp>(
-        loc, r, constantOne(rewriter, loc, vtp),
-        constantIndex(rewriter, loc, 0));
+    return builder.create<vector::InsertElementOp>(
+        loc, r, constantOne(builder, loc, vtp), constantIndex(builder, loc, 0));
   case kAnd:
   case kOr:
     // Initialize reduction vector to: | r | .. | r | r |
-    return rewriter.create<vector::BroadcastOp>(loc, vtp, r);
+    return builder.create<vector::BroadcastOp>(loc, vtp, r);
   }
   llvm_unreachable("unknown reduction kind");
 }
 
 /// Generates final value for a vector reduction.
-static Value genVectorReducEnd(CodeGen &codegen, PatternRewriter &rewriter,
+static Value genVectorReducEnd(CodeGen &codegen, OpBuilder &builder,
                                Location loc, VectorType vtp) {
   vector::CombiningKind kind = getCombiningKind(codegen.redKind);
-  return rewriter.create<vector::ReductionOp>(loc, kind, codegen.redVal);
+  return builder.create<vector::ReductionOp>(loc, kind, codegen.redVal);
 }
 
 /// Updates scalarized reduction value.
 static void updateReduc(Merger &merger, CodeGen &codegen, Value reduc) {
   assert(codegen.redKind != kNoReduc);
   codegen.redVal = merger.exp(codegen.redExp).val = reduc;
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (statements and expressions).
 //===----------------------------------------------------------------------===//
 
 /// Generates buffer for the output tensor. Note that all sparse kernels
 /// assume that when all elements are written to (viz. x(i) = y(i) * z(i)),
 /// the output buffer is already initialized to all zeroes and only nonzeroes
 /// values are computed and written out. For updates (viz. x(i) += y(i) * z(i)),
 /// only nonzeroes values are used for the updates and no assumption on the
 /// original contents of the output buffer is necessary..
-static Value genOutputBuffer(CodeGen &codegen, PatternRewriter &rewriter,
+static Value genOutputBuffer(CodeGen &codegen, OpBuilder &builder,
                              linalg::GenericOp op, MemRefType denseTp,
                              ArrayRef<Value> args) {
   Location loc = op.getLoc();
   Value tensor = op.getOutputOperand(0)->get();
   // The output tensor simply could materialize from the buffer that will
   // be generated for the tensor present in the outs() clause. This has
   // the major advantage that the sparse kernel only updates the nonzero
   // positions for the output tensor.
   if (isInPlace(tensor))
-    return rewriter.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
+    return builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
   // By default, a new buffer is allocated which is initialized to the
   // tensor defined in the outs() clause. This is always correct but
   // introduces a dense initialization component that may negatively
   // impact the running complexity of the sparse kernel. If the tensor
   // materializes into the computation, we need to preserve the zero
   // initialization assumption of all sparse output buffers.
-  Value alloc = rewriter.create<memref::AllocOp>(loc, denseTp, args);
+  Value alloc = builder.create<memref::AllocOp>(loc, denseTp, args);
   if (isMaterializing(tensor)) {
-    Value zero = constantZero(rewriter, loc, denseTp.getElementType());
-    rewriter.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{alloc});
+    Value zero = constantZero(builder, loc, denseTp.getElementType());
+    builder.create<linalg::FillOp>(loc, ValueRange{zero}, ValueRange{alloc});
   } else {
     Value init =
-        rewriter.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
-    rewriter.create<memref::CopyOp>(loc, init, alloc);
+        builder.create<bufferization::ToMemrefOp>(loc, denseTp, tensor);
+    builder.create<memref::CopyOp>(loc, init, alloc);
   }
   return alloc;
 }
 
 /// Local bufferization of all dense and sparse data structures.
 /// This code enables testing the first prototype sparse compiler.
 // TODO: replace this with a proliferated bufferization strategy
-static void genBuffers(Merger &merger, CodeGen &codegen,
-                       PatternRewriter &rewriter, linalg::GenericOp op) {
+static void genBuffers(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                       linalg::GenericOp op) {
   Location loc = op.getLoc();
   assert(op.getNumInputsAndOutputs() == op.getNumInputs() + 1);
   // For every tensor, find lower and upper bound on dimensions, set the
   // same bounds on loop indices, and obtain dense or sparse buffer(s).
   SmallVector<Value, 4> args;
   for (OpOperand *t : op.getInputAndOutputOperands()) {
     unsigned tensor = t->getOperandNumber();
     auto shape = op.getShape(t);
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
     // Scan all dimensions of current tensor.
     args.clear();
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       if (a.getKind() != AffineExprKind::DimId)
         continue; // compound
       unsigned idx = a.cast<AffineDimExpr>().getPosition();
       // Handle sparse storage schemes.
       if (merger.isDim(tensor, idx, Dim::kSparse)) {
         auto dynShape = {ShapedType::kDynamicSize};
         auto ptrTp =
-            MemRefType::get(dynShape, getPointerOverheadType(rewriter, enc));
+            MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
-            MemRefType::get(dynShape, getIndexOverheadType(rewriter, enc));
-        Value dim = constantIndex(rewriter, loc, d);
+            MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
+        Value dim = constantIndex(builder, loc, d);
         // Generate sparse primitives to obtains pointer and indices.
         codegen.pointers[tensor][idx] =
-            rewriter.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
+            builder.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
         codegen.indices[tensor][idx] =
-            rewriter.create<ToIndicesOp>(loc, indTp, t->get(), dim);
+            builder.create<ToIndicesOp>(loc, indTp, t->get(), dim);
       }
       // Find upper bound in current dimension.
       unsigned p = perm(enc, d);
-      Value up = linalg::createOrFoldDimOp(rewriter, loc, t->get(), p);
+      Value up = linalg::createOrFoldDimOp(builder, loc, t->get(), p);
       if (ShapedType::isDynamic(shape[p]))
         args.push_back(up);
       assert(codegen.highs[tensor][idx] == nullptr);
       codegen.sizes[idx] = codegen.highs[tensor][idx] = up;
     }
     // Perform the required bufferization. Dense inputs materialize
     // from the input tensors. Dense outputs need special handling.
     // Sparse inputs use sparse primitives to obtain the values.
     // We also accept in-place all-dense annotated "sparse" outputs.
     Type elementType = getElementTypeOrSelf(t->get().getType());
     if (!enc) {
       // Non-annotated dense tensors.
       auto denseTp = MemRefType::get(shape, elementType);
       if (tensor < op.getNumInputs())
         codegen.buffers[tensor] =
-            rewriter.create<bufferization::ToMemrefOp>(loc, denseTp, t->get());
+            builder.create<bufferization::ToMemrefOp>(loc, denseTp, t->get());
       else
         codegen.buffers[tensor] =
-            genOutputBuffer(codegen, rewriter, op, denseTp, args);
+            genOutputBuffer(codegen, builder, op, denseTp, args);
     } else if (t == codegen.sparseOut) {
       // True sparse output needs a lexIdx array.
-      Value rank = constantIndex(rewriter, loc, op.getRank(t));
+      Value rank = constantIndex(builder, loc, op.getRank(t));
       auto dynShape = {ShapedType::kDynamicSize};
-      auto memTp = MemRefType::get(dynShape, rewriter.getIndexType());
-      codegen.lexIdx = rewriter.create<memref::AllocaOp>(loc, memTp, rank);
+      auto memTp = MemRefType::get(dynShape, builder.getIndexType());
+      codegen.lexIdx = builder.create<memref::AllocaOp>(loc, memTp, rank);
     } else {
       // Annotated sparse tensors.
       auto dynShape = {ShapedType::kDynamicSize};
       auto sparseTp = MemRefType::get(dynShape, elementType);
       codegen.buffers[tensor] =
-          rewriter.create<ToValuesOp>(loc, sparseTp, t->get());
+          builder.create<ToValuesOp>(loc, sparseTp, t->get());
     }
   }
 }
 
 /// Constructs vector type.
 static VectorType vectorType(CodeGen &codegen, Type etp) {
   unsigned numScalableDims = codegen.options.enableVLAVectorization;
   return VectorType::get(codegen.curVecLength, etp, numScalableDims);
 }
 
 /// Constructs vector type from pointer.
 static VectorType vectorType(CodeGen &codegen, Value ptr) {
   return vectorType(codegen, ptr.getType().cast<MemRefType>().getElementType());
 }
 
 /// Constructs vector iteration mask.
-static Value genVectorMask(CodeGen &codegen, PatternRewriter &rewriter,
-                           Value iv, Value lo, Value hi, Value step) {
+static Value genVectorMask(CodeGen &codegen, OpBuilder &builder, Value iv,
+                           Value lo, Value hi, Value step) {
   Location loc = iv.getLoc();
-  VectorType mtp = vectorType(codegen, rewriter.getI1Type());
+  VectorType mtp = vectorType(codegen, builder.getI1Type());
   // Special case if the vector length evenly divides the trip count (for
   // example, "for i = 0, 128, 16"). A constant all-true mask is generated
   // so that all subsequent masked memory operations are immediately folded
   // into unconditional memory operations.
   IntegerAttr loInt, hiInt, stepInt;
   if (matchPattern(lo, m_Constant(&loInt)) &&
       matchPattern(hi, m_Constant(&hiInt)) &&
       matchPattern(step, m_Constant(&stepInt))) {
     if (((hiInt.getInt() - loInt.getInt()) % stepInt.getInt()) == 0)
-      return rewriter.create<vector::BroadcastOp>(
-          loc, mtp, constantI1(rewriter, loc, true));
+      return builder.create<vector::BroadcastOp>(
+          loc, mtp, constantI1(builder, loc, true));
   }
   // Otherwise, generate a vector mask that avoids overrunning the upperbound
   // during vector execution. Here we rely on subsequent loop optimizations to
   // avoid executing the mask in all iterations, for example, by splitting the
   // loop into an unconditional vector loop and a scalar cleanup loop.
   auto minMap = AffineMap::get(
       /*dimCount=*/2, /*symbolCount=*/1,
-      {rewriter.getAffineSymbolExpr(0),
-       rewriter.getAffineDimExpr(0) - rewriter.getAffineDimExpr(1)},
-      rewriter.getContext());
+      {builder.getAffineSymbolExpr(0),
+       builder.getAffineDimExpr(0) - builder.getAffineDimExpr(1)},
+      builder.getContext());
   Value end =
-      rewriter.createOrFold<AffineMinOp>(loc, minMap, ValueRange{hi, iv, step});
-  return rewriter.create<vector::CreateMaskOp>(loc, mtp, end);
+      builder.createOrFold<AffineMinOp>(loc, minMap, ValueRange{hi, iv, step});
+  return builder.create<vector::CreateMaskOp>(loc, mtp, end);
 }
 
 /// Generates a vectorized load lhs = a[ind[lo:hi]] or lhs = a[lo:hi].
-static Value genVectorLoad(CodeGen &codegen, PatternRewriter &rewriter,
-                           Value ptr, ArrayRef<Value> args) {
+static Value genVectorLoad(CodeGen &codegen, OpBuilder &builder, Value ptr,
+                           ArrayRef<Value> args) {
   Location loc = ptr.getLoc();
   VectorType vtp = vectorType(codegen, ptr);
-  Value pass = constantZero(rewriter, loc, vtp);
+  Value pass = constantZero(builder, loc, vtp);
   if (args.back().getType().isa<VectorType>()) {
     SmallVector<Value, 4> scalarArgs(args.begin(), args.end());
     Value indexVec = args.back();
-    scalarArgs.back() = constantIndex(rewriter, loc, 0);
-    return rewriter.create<vector::GatherOp>(
-        loc, vtp, ptr, scalarArgs, indexVec, codegen.curVecMask, pass);
+    scalarArgs.back() = constantIndex(builder, loc, 0);
+    return builder.create<vector::GatherOp>(loc, vtp, ptr, scalarArgs, indexVec,
+                                            codegen.curVecMask, pass);
   }
-  return rewriter.create<vector::MaskedLoadOp>(loc, vtp, ptr, args,
-                                               codegen.curVecMask, pass);
+  return builder.create<vector::MaskedLoadOp>(loc, vtp, ptr, args,
 }
 
 /// Generates a vectorized store a[ind[lo:hi]] = rhs or a[lo:hi] = rhs.
-static void genVectorStore(CodeGen &codegen, PatternRewriter &rewriter,
-                           Value rhs, Value ptr, ArrayRef<Value> args) {
+static void genVectorStore(CodeGen &codegen, OpBuilder &builder, Value rhs,
+                           Value ptr, ArrayRef<Value> args) {
   Location loc = ptr.getLoc();
   if (args.back().getType().isa<VectorType>()) {
     SmallVector<Value, 4> scalarArgs(args.begin(), args.end());
     Value indexVec = args.back();
-    scalarArgs.back() = constantIndex(rewriter, loc, 0);
-    rewriter.create<vector::ScatterOp>(loc, ptr, scalarArgs, indexVec,
-                                       codegen.curVecMask, rhs);
+    scalarArgs.back() = constantIndex(builder, loc, 0);
+    builder.create<vector::ScatterOp>(loc, ptr, scalarArgs, indexVec,
     return;
   }
-  rewriter.create<vector::MaskedStoreOp>(loc, ptr, args, codegen.curVecMask,
-                                         rhs);
+  builder.create<vector::MaskedStoreOp>(loc, ptr, args, codegen.curVecMask,
 }
 
 /// Generates a vectorized invariant. Here we rely on subsequent loop
 /// optimizations to hoist the invariant broadcast out of the vector loop.
-static Value genVectorInvariantValue(CodeGen &codegen,
-                                     PatternRewriter &rewriter, Value val) {
+static Value genVectorInvariantValue(CodeGen &codegen, OpBuilder &builder,
+                                     Value val) {
   VectorType vtp = vectorType(codegen, val.getType());
-  return rewriter.create<vector::BroadcastOp>(val.getLoc(), vtp, val);
+  return builder.create<vector::BroadcastOp>(val.getLoc(), vtp, val);
 }
 
 /// Generates an affine expression.
 //
 // TODO: generalize for sparse tensor subscripts
 //
-static Value genAffine(CodeGen &codegen, PatternRewriter &rewriter,
-                       AffineExpr a, Location loc) {
+static Value genAffine(CodeGen &codegen, OpBuilder &builder, AffineExpr a,
+                       Location loc) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
     return codegen.loops[idx]; // universal dense index
   }
   case AffineExprKind::Add: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return rewriter.create<arith::AddIOp>(
-        loc, genAffine(codegen, rewriter, binOp.getLHS(), loc),
-        genAffine(codegen, rewriter, binOp.getRHS(), loc));
+    return builder.create<arith::AddIOp>(
+        loc, genAffine(codegen, builder, binOp.getLHS(), loc),
+        genAffine(codegen, builder, binOp.getRHS(), loc));
   }
   case AffineExprKind::Mul: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return rewriter.create<arith::MulIOp>(
-        loc, genAffine(codegen, rewriter, binOp.getLHS(), loc),
-        genAffine(codegen, rewriter, binOp.getRHS(), loc));
+    return builder.create<arith::MulIOp>(
+        loc, genAffine(codegen, builder, binOp.getLHS(), loc),
+        genAffine(codegen, builder, binOp.getRHS(), loc));
   }
   case AffineExprKind::Constant: {
     int64_t c = a.cast<AffineConstantExpr>().getValue();
-    return constantIndex(rewriter, loc, c);
+    return constantIndex(builder, loc, c);
   }
   default:
     llvm_unreachable("unexpected affine subscript");
   }
 }
 
 /// Generates index for load/store on sparse tensor.
 static Value genIndex(CodeGen &codegen, linalg::GenericOp op, OpOperand *t) {
   auto map = op.getTiedIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   AffineExpr a = map.getResult(perm(enc, map.getNumResults() - 1));
   assert(a.getKind() == AffineExprKind::DimId);
   unsigned idx = a.cast<AffineDimExpr>().getPosition();
   return codegen.loops[idx];
 }
 
 /// Generates subscript for load/store on a dense or sparse tensor.
-static Value genSubscript(CodeGen &codegen, PatternRewriter &rewriter,
+static Value genSubscript(CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, OpOperand *t,
                           SmallVector<Value, 4> &args) {
   unsigned tensor = t->getOperandNumber();
   auto map = op.getTiedIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   unsigned rank = map.getNumResults();
   if (enc) {
     // Note that currently, all sparse subscripts are simple.
     // TODO: accept affine too?
     AffineExpr a = map.getResult(perm(enc, rank - 1));
     assert(a.getKind() == AffineExprKind::DimId);
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
     assert(codegen.pidxs[tensor][idx] != nullptr);
     args.push_back(codegen.pidxs[tensor][idx]); // position index
   } else {
     for (unsigned d = 0; d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
-      args.push_back(genAffine(codegen, rewriter, a, op.getLoc()));
+      args.push_back(genAffine(codegen, builder, a, op.getLoc()));
     }
   }
   return codegen.buffers[tensor];
 }
 
 /// Generates insertion code to implement dynamic tensor load.
-static Value genInsertionLoad(CodeGen &codegen, PatternRewriter &rewriter,
+static Value genInsertionLoad(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t) {
   Location loc = op.getLoc();
   // Direct lexicographic index order, tensor loads as zero.
   if (!codegen.expValues) {
     Type tp = getElementTypeOrSelf(t->get().getType());
-    return constantZero(rewriter, loc, tp);
+    return constantZero(builder, loc, tp);
   }
   // Load from expanded access pattern.
   Value index = genIndex(codegen, op, t);
-  return rewriter.create<memref::LoadOp>(loc, codegen.expValues, index);
+  return builder.create<memref::LoadOp>(loc, codegen.expValues, index);
 }
 
 /// Generates insertion code to implement dynamic tensor store.
-static void genInsertionStore(CodeGen &codegen, PatternRewriter &rewriter,
+static void genInsertionStore(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t, Value rhs) {
   Location loc = op.getLoc();
   // Direct insertion in lexicographic index order.
   if (!codegen.expValues) {
-    rewriter.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
+    builder.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, rhs);
     return;
   }
   // Generates insertion code along expanded access pattern.
   //   if (!expFilled[i]) then
   //     expFilled[i] = true
   //     expAdded[inserts++] = i
   //   endif
   //   values[i] = rhs
   Value index = genIndex(codegen, op, t);
-  Value fval = constantI1(rewriter, loc, false);
-  Value tval = constantI1(rewriter, loc, true);
+  Value fval = constantI1(builder, loc, false);
+  Value tval = constantI1(builder, loc, true);
   // If statement.
-  Value filled = rewriter.create<memref::LoadOp>(loc, codegen.expFilled, index);
-  Value cond = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                              filled, fval);
-  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, rewriter.getIndexType(),
-                                              cond, /*else=*/true);
+  Value filled = builder.create<memref::LoadOp>(loc, codegen.expFilled, index);
+  Value cond = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, builder.getIndexType(), cond,
+                                             /*else=*/true);
   // True branch.
-  rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
-  rewriter.create<memref::StoreOp>(loc, tval, codegen.expFilled, index);
-  rewriter.create<memref::StoreOp>(loc, index, codegen.expAdded,
-                                   codegen.expCount);
-  Value one = constantIndex(rewriter, loc, 1);
-  Value add = rewriter.create<arith::AddIOp>(loc, codegen.expCount, one);
-  rewriter.create<scf::YieldOp>(loc, add);
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+  builder.create<memref::StoreOp>(loc, tval, codegen.expFilled, index);
+  builder.create<memref::StoreOp>(loc, index, codegen.expAdded,
+  Value one = constantIndex(builder, loc, 1);
+  Value add = builder.create<arith::AddIOp>(loc, codegen.expCount, one);
+  builder.create<scf::YieldOp>(loc, add);
   // False branch.
-  rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
-  rewriter.create<scf::YieldOp>(loc, codegen.expCount);
-  rewriter.setInsertionPointAfter(ifOp);
+  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+  builder.create<scf::YieldOp>(loc, codegen.expCount);
+  builder.setInsertionPointAfter(ifOp);
   // Value assignment.
   codegen.expCount = ifOp.getResult(0);
-  rewriter.create<memref::StoreOp>(loc, rhs, codegen.expValues, index);
+  builder.create<memref::StoreOp>(loc, rhs, codegen.expValues, index);
 }
 
 /// Generates a load on a dense or sparse tensor.
-static Value genTensorLoad(Merger &merger, CodeGen &codegen,
-                           PatternRewriter &rewriter, linalg::GenericOp op,
-                           unsigned exp) {
+static Value genTensorLoad(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                           linalg::GenericOp op, unsigned exp) {
   // Test if the load was hoisted to a higher loop nest.
   Value val = merger.exp(exp).val;
   if (val) {
     if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
-      return genVectorInvariantValue(codegen, rewriter, val);
+      return genVectorInvariantValue(codegen, builder, val);
     return val;
   }
   // Load during insertion.
   OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
   if (t == codegen.sparseOut)
-    return genInsertionLoad(codegen, rewriter, op, t);
+    return genInsertionLoad(codegen, builder, op, t);
   // Actual load.
   SmallVector<Value, 4> args;
-  Value ptr = genSubscript(codegen, rewriter, op, t, args);
+  Value ptr = genSubscript(codegen, builder, op, t, args);
   if (codegen.curVecLength > 1)
-    return genVectorLoad(codegen, rewriter, ptr, args);
-  return rewriter.create<memref::LoadOp>(op.getLoc(), ptr, args);
+    return genVectorLoad(codegen, builder, ptr, args);
+  return builder.create<memref::LoadOp>(op.getLoc(), ptr, args);
 }
 
 /// Generates a store on a dense or sparse tensor.
-static void genTensorStore(Merger &merger, CodeGen &codegen,
-                           PatternRewriter &rewriter, linalg::GenericOp op,
-                           unsigned exp, Value rhs) {
+static void genTensorStore(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                           linalg::GenericOp op, unsigned exp, Value rhs) {
   Location loc = op.getLoc();
   // Test if this is a scalarized reduction.
   if (codegen.redVal) {
     if (codegen.curVecLength > 1)
-      rhs = rewriter.create<arith::SelectOp>(loc, codegen.curVecMask, rhs,
-                                             codegen.redVal);
+      rhs = builder.create<arith::SelectOp>(loc, codegen.curVecMask, rhs,
     updateReduc(merger, codegen, rhs);
     return;
   }
   // Store during insertion.
   OpOperand *t = op.getOutputOperand(0);
   if (t == codegen.sparseOut) {
     if (!rhs) {
       // Only unary and binary are allowed to return uninitialized rhs
       // to indicate missing output.
       assert(merger.exp(exp).kind == kUnary || merger.exp(exp).kind == kBinary);
     } else {
-      genInsertionStore(codegen, rewriter, op, t, rhs);
+      genInsertionStore(codegen, builder, op, t, rhs);
     }
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
-  Value ptr = genSubscript(codegen, rewriter, op, t, args);
+  Value ptr = genSubscript(codegen, builder, op, t, args);
   if (codegen.curVecLength > 1)
-    genVectorStore(codegen, rewriter, rhs, ptr, args);
+    genVectorStore(codegen, builder, rhs, ptr, args);
   else
-    rewriter.create<memref::StoreOp>(loc, rhs, ptr, args);
+    builder.create<memref::StoreOp>(loc, rhs, ptr, args);
 }
 
 /// Generates a pointer/index load from the sparse storage scheme. Narrower
 /// data types need to be zero extended before casting the value into the
 /// index type used for looping and indexing.
-static Value genLoad(CodeGen &codegen, PatternRewriter &rewriter, Location loc,
+static Value genLoad(CodeGen &codegen, OpBuilder &builder, Location loc,
                      Value ptr, Value s) {
   // See https://llvm.org/docs/GetElementPtr.html for some background on
   // the complications described below.
   if (codegen.curVecLength > 1) {
     // Since the index vector is used in a subsequent gather/scatter operations,
     // which effectively defines an unsigned pointer + signed index, we must
     // zero extend the vector to an index width. For 8-bit and 16-bit values,
     // an 32-bit index width suffices. For 32-bit values, zero extending the
     // elements into 64-bit loses some performance since the 32-bit indexed
     // gather/scatter is more efficient than the 64-bit index variant (if the
     // negative 32-bit index space is unused, the enableSIMDIndex32 flag can
     // preserve this performance). For 64-bit values, there is no good way
     // to state that the indices are unsigned, with creates the potential of
     // incorrect address calculations in the unlikely case we need such
     // extremely large offsets.
     Type etp = ptr.getType().cast<MemRefType>().getElementType();
-    Value vload = genVectorLoad(codegen, rewriter, ptr, {s});
+    Value vload = genVectorLoad(codegen, builder, ptr, {s});
     if (!etp.isa<IndexType>()) {
       if (etp.getIntOrFloatBitWidth() < 32)
-        vload = rewriter.create<arith::ExtUIOp>(
-            loc, vectorType(codegen, rewriter.getI32Type()), vload);
+        vload = builder.create<arith::ExtUIOp>(
+            loc, vectorType(codegen, builder.getI32Type()), vload);
       else if (etp.getIntOrFloatBitWidth() < 64 &&
                !codegen.options.enableSIMDIndex32)
-        vload = rewriter.create<arith::ExtUIOp>(
-            loc, vectorType(codegen, rewriter.getI64Type()), vload);
+        vload = builder.create<arith::ExtUIOp>(
+            loc, vectorType(codegen, builder.getI64Type()), vload);
     }
     return vload;
   }
   // For the scalar case, we simply zero extend narrower indices into 64-bit
   // values before casting to index without a performance penalty. Here too,
   // however, indices that already are 64-bit, in theory, cannot express the
   // full range as explained above.
-  Value load = rewriter.create<memref::LoadOp>(loc, ptr, s);
+  Value load = builder.create<memref::LoadOp>(loc, ptr, s);
   if (!load.getType().isa<IndexType>()) {
     if (load.getType().getIntOrFloatBitWidth() < 64)
-      load = rewriter.create<arith::ExtUIOp>(loc, rewriter.getI64Type(), load);
+      load = builder.create<arith::ExtUIOp>(loc, builder.getI64Type(), load);
     load =
-        rewriter.create<arith::IndexCastOp>(loc, rewriter.getIndexType(), load);
+        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), load);
   }
   return load;
 }
 
 /// Generates an invariant value.
 static Value genInvariantValue(Merger &merger, CodeGen &codegen,
-                               PatternRewriter &rewriter, unsigned exp) {
+                               OpBuilder &builder, unsigned exp) {
   Value val = merger.exp(exp).val;
   if (codegen.curVecLength > 1)
-    return genVectorInvariantValue(codegen, rewriter, val);
+    return genVectorInvariantValue(codegen, builder, val);
   return val;
 }
 
 /// Generates an address computation "sz * p + i".
-static Value genAddress(CodeGen &codegen, PatternRewriter &rewriter,
-                        Location loc, Value size, Value p, Value i) {
-  Value mul = rewriter.create<arith::MulIOp>(loc, size, p);
+static Value genAddress(CodeGen &codegen, OpBuilder &builder, Location loc,
+                        Value size, Value p, Value i) {
+  Value mul = builder.create<arith::MulIOp>(loc, size, p);
   if (auto vtp = i.getType().dyn_cast<VectorType>()) {
     Value inv =
-        rewriter.create<arith::IndexCastOp>(loc, vtp.getElementType(), mul);
-    mul = genVectorInvariantValue(codegen, rewriter, inv);
+        builder.create<arith::IndexCastOp>(loc, vtp.getElementType(), mul);
+    mul = genVectorInvariantValue(codegen, builder, inv);
   }
-  return rewriter.create<arith::AddIOp>(loc, mul, i);
+  return builder.create<arith::AddIOp>(loc, mul, i);
 }
 
 /// Generates an index value.
-static Value genIndexValue(Merger &merger, CodeGen &codegen,
-                           PatternRewriter &rewriter, unsigned exp,
-                           unsigned ldx) {
+static Value genIndexValue(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                           unsigned exp, unsigned ldx) {
   unsigned idx = merger.exp(exp).index;
   Value ival = codegen.loops[idx];
   Type itype = ival.getType();
   // During vectorization, we either encounter:
   // (1) indices already in vector form, as in ... = ind[lo:hi], good to go, or
   // (2) single index, as in ... = i, must convert to [i, i+1, ...] for inner i.
   unsigned vl = codegen.curVecLength;
   if (vl > 1 && !itype.isa<VectorType>()) {
     Location loc = ival.getLoc();
     VectorType vtp = vectorType(codegen, itype);
-    ival = rewriter.create<vector::BroadcastOp>(loc, vtp, ival);
+    ival = builder.create<vector::BroadcastOp>(loc, vtp, ival);
     if (idx == ldx) {
       Value incr;
       if (vtp.isScalable()) {
-        Type stepvty = vectorType(codegen, rewriter.getI64Type());
-        Value stepv = rewriter.create<LLVM::StepVectorOp>(loc, stepvty);
-        incr = rewriter.create<arith::IndexCastOp>(loc, vtp, stepv);
+        Type stepvty = vectorType(codegen, builder.getI64Type());
+        Value stepv = builder.create<LLVM::StepVectorOp>(loc, stepvty);
+        incr = builder.create<arith::IndexCastOp>(loc, vtp, stepv);
       } else {
         SmallVector<APInt, 4> integers;
         for (unsigned i = 0; i < vl; i++)
           integers.push_back(APInt(/*width=*/64, i));
         auto values = DenseElementsAttr::get(vtp, integers);
-        incr = rewriter.create<arith::ConstantOp>(loc, vtp, values);
+        incr = builder.create<arith::ConstantOp>(loc, vtp, values);
       }
-      ival = rewriter.create<arith::AddIOp>(loc, ival, incr);
+      ival = builder.create<arith::AddIOp>(loc, ival, incr);
     }
   }
   return ival;
 }
 
 /// Recursively generates tensor expression.
-static Value genExp(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
+static Value genExp(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                     linalg::GenericOp op, unsigned exp, unsigned ldx) {
   Location loc = op.getLoc();
   if (exp == -1u)
     return Value();
   if (merger.exp(exp).kind == Kind::kTensor)
     return genTensorLoad(merger, codegen, rewriter, op, exp);
   if (merger.exp(exp).kind == Kind::kInvariant)
     return genInvariantValue(merger, codegen, rewriter, exp);
return isInvariantAffine(codegen, binOp.getLHS(), ldx, atLevel) &&
            isInvariantAffine(codegen, binOp.getRHS(), ldx, atLevel);
   }
   default:
     return true;
   }
 }
 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
-static void genInvariants(Merger &merger, CodeGen &codegen,
-                          PatternRewriter &rewriter, linalg::GenericOp op,
-                          unsigned exp, unsigned ldx, bool atStart,
-                          Kind last = Kind::kTensor) {
+static void genInvariants(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                          linalg::GenericOp op, unsigned exp, unsigned ldx,
+                          bool atStart, Kind last = Kind::kTensor) {
   if (exp == -1u)
     return;
   if (merger.exp(exp).kind == Kind::kTensor) {
     // Inspect tensor indices.
     bool atLevel = ldx == -1u;
     OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       if (!isInvariantAffine(codegen, a, ldx, atLevel))
         return; // still in play
     }
     // All exhausted at this level (atLevel denotes exactly at this level).
     if (!atLevel)
       return;
     OpOperand *lhs = op.getOutputOperand(0);
     if (lhs == t) {
       // Start or end a scalarized reduction
       if (atStart) {
-        Value load = genTensorLoad(merger, codegen, rewriter, op, exp);
+        Value load = genTensorLoad(merger, codegen, builder, op, exp);
         codegen.redKind = getReduction(last);
         codegen.redExp = exp;
         updateReduc(merger, codegen, load);
       } else {
         Value redVal = codegen.redVal;
         updateReduc(merger, codegen, Value());
         codegen.redExp = -1u;
         codegen.redKind = kNoReduc;
-        genTensorStore(merger, codegen, rewriter, op, exp, redVal);
+        genTensorStore(merger, codegen, builder, op, exp, redVal);
       }
     } else {
       // Start or end loop invariant hoisting of a tensor load.
       merger.exp(exp).val =
-          atStart ? genTensorLoad(merger, codegen, rewriter, op, exp) : Value();
+          atStart ? genTensorLoad(merger, codegen, builder, op, exp) : Value();
     }
   } else if (merger.exp(exp).kind != Kind::kInvariant &&
              merger.exp(exp).kind != Kind::kIndex) {
     // Traverse into the binary operations. Note that we only hoist
     // tensor loads, since subsequent MLIR/LLVM passes know how to
     // deal with all other kinds of derived loop invariants.
     Kind last = merger.exp(exp).kind;
     unsigned e0 = merger.exp(exp).children.e0;
     unsigned e1 = merger.exp(exp).children.e1;
-    genInvariants(merger, codegen, rewriter, op, e0, ldx, atStart, last);
-    genInvariants(merger, codegen, rewriter, op, e1, ldx, atStart, last);
+    genInvariants(merger, codegen, builder, op, e0, ldx, atStart, last);
+    genInvariants(merger, codegen, builder, op, e1, ldx, atStart, last);
   }
 }
 
 /// Generates an expanded access pattern in innermost dimension.
-static void genExpansion(Merger &merger, CodeGen &codegen,
-                         PatternRewriter &rewriter, linalg::GenericOp op,
-                         unsigned at, bool atStart) {
+static void genExpansion(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                         linalg::GenericOp op, unsigned at, bool atStart) {
   OpOperand *lhs = codegen.sparseOut;
   if (!lhs || codegen.outerParNest != op.getRank(lhs) - 1 ||
       at != codegen.outerParNest)
     return; // not needed at this level
   // Generate start or end of an expanded access pattern.
   Value tensor = lhs->get();
   Location loc = op.getLoc();
   if (atStart) {
     auto dynShape = {ShapedType::kDynamicSize};
     Type etp = tensor.getType().cast<ShapedType>().getElementType();
     Type t1 = MemRefType::get(dynShape, etp);
-    Type t2 = MemRefType::get(dynShape, rewriter.getI1Type());
-    Type t3 = MemRefType::get(dynShape, rewriter.getIndexType());
-    Type t4 = rewriter.getIndexType();
+    Type t2 = MemRefType::get(dynShape, builder.getI1Type());
+    Type t3 = MemRefType::get(dynShape, builder.getIndexType());
+    Type t4 = builder.getIndexType();
     auto res =
-        rewriter.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
+        builder.create<ExpandOp>(loc, TypeRange({t1, t2, t3, t4}), tensor);
     assert(res.getNumResults() == 4);
     assert(!codegen.expValues);
     codegen.expValues = res.getResult(0);
     codegen.expFilled = res.getResult(1);
     codegen.expAdded = res.getResult(2);
     codegen.expCount = res.getResult(3);
   } else {
     assert(codegen.expValues);
-    rewriter.create<CompressOp>(loc, tensor, codegen.lexIdx, codegen.expValues,
-                                codegen.expFilled, codegen.expAdded,
-                                codegen.expCount);
+    builder.create<CompressOp>(loc, tensor, codegen.lexIdx, codegen.expValues,
     codegen.expValues = codegen.expFilled = codegen.expAdded =
         codegen.expCount = Value();
   }
 }
 
 /// Generates initialization code for the subsequent loop sequence at
 /// current index level. Returns true if the loop sequence needs to
 /// maintain the universal index.
-static bool genInit(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
+static bool genInit(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned at, BitVector &inits) {
   bool needsUniv = false;
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse positions.
   for (unsigned b = 0, be = inits.size(); b < be; b++) {
     if (inits[b]) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       if (merger.isDim(b, Dim::kSparse)) {
         // Initialize sparse index.
         unsigned pat = at;
         for (; pat != 0; pat--) {
           if (codegen.pidxs[tensor][topSort[pat - 1]])
             break;
         }
         Value ptr = codegen.pointers[tensor][idx];
-        Value one = constantIndex(rewriter, loc, 1);
-        Value p0 = (pat == 0) ? constantIndex(rewriter, loc, 0)
+        Value one = constantIndex(builder, loc, 1);
+        Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
                               : codegen.pidxs[tensor][topSort[pat - 1]];
-        codegen.pidxs[tensor][idx] = genLoad(codegen, rewriter, loc, ptr, p0);
-        Value p1 = rewriter.create<arith::AddIOp>(loc, p0, one);
-        codegen.highs[tensor][idx] = genLoad(codegen, rewriter, loc, ptr, p1);
+        codegen.pidxs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p0);
+        Value p1 = builder.create<arith::AddIOp>(loc, p0, one);
+        codegen.highs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p1);
       } else {
         // Dense index still in play.
         needsUniv = true;
       }
     }
   }
 
   // Initialize the universal dense index.
-  codegen.loops[idx] = constantIndex(rewriter, loc, 0);
+  codegen.loops[idx] = constantIndex(builder, loc, 0);
   return needsUniv;
 }
 
 /// Returns vectorization strategy. Any implicit inner loop in the Linalg
 /// operation is a candidate. Whether it is actually converted to SIMD code
 /// depends on the requested strategy.
 static bool isVectorFor(CodeGen &codegen, bool isInner, bool isReduction,
                         bool isSparse) {
if (!getSparseTensorEncoding(t->get().getType())) {
           return false;
       }
     }
   }
   return true;
 }
 
 /// Generates a for-loop on a single index.
-static Operation *genFor(Merger &merger, CodeGen &codegen,
-                         PatternRewriter &rewriter, linalg::GenericOp op,
-                         bool isOuter, bool isInner, unsigned idx,
-                         BitVector &indices) {
+static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                         linalg::GenericOp op, bool isOuter, bool isInner,
+                         unsigned idx, BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
   auto iteratorTypes = op.iterator_types().getValue();
   bool isReduction = isReductionIterator(iteratorTypes[idx]);
   bool isSparse = merger.isDim(fb, Dim::kSparse);
   bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
   // Prepare vector length.
   if (isVector)
     codegen.curVecLength = codegen.options.vectorLength;
 
   // Loop bounds and increment.
   Location loc = op.getLoc();
   Value lo = isSparse ? codegen.pidxs[tensor][idx] : codegen.loops[idx];
   Value hi = isSparse ? codegen.highs[tensor][idx] : codegen.sizes[idx];
-  Value step = constantIndex(rewriter, loc, codegen.curVecLength);
+  Value step = constantIndex(builder, loc, codegen.curVecLength);
   if (isVector && codegen.options.enableVLAVectorization) {
-    Value vscale = rewriter.create<vector::VectorScaleOp>(
-        loc, IndexType::get(rewriter.getContext()));
-    step = rewriter.create<arith::MulIOp>(loc, vscale, step);
+    Value vscale = builder.create<vector::VectorScaleOp>(
+        loc, IndexType::get(builder.getContext()));
+    step = builder.create<arith::MulIOp>(loc, vscale, step);
   }
 
   // Emit a parallel loop.
   if (isParallel) {
     assert(!isVector);
-    scf::ParallelOp parOp = rewriter.create<scf::ParallelOp>(loc, lo, hi, step);
+    scf::ParallelOp parOp = builder.create<scf::ParallelOp>(loc, lo, hi, step);
     if (isSparse)
       codegen.pidxs[tensor][idx] = parOp.getInductionVars()[0];
     else
       codegen.loops[idx] = parOp.getInductionVars()[0];
-    rewriter.setInsertionPointToStart(parOp.getBody());
+    builder.setInsertionPointToStart(parOp.getBody());
     return parOp;
   }
 
   // Emit a sequential or vector loop.
   SmallVector<Value, 4> operands;
   if (codegen.redVal) {
     // In a vector loop, bring reduction into SIMD form, if not already.
     if (isVector && !codegen.redVal.getType().isa<VectorType>()) {
       VectorType vtp = vectorType(codegen, codegen.redVal.getType());
-      Value vred = genVectorReducInit(codegen, rewriter, loc, vtp);
+      Value vred = genVectorReducInit(codegen, builder, loc, vtp);
       updateReduc(merger, codegen, vred);
     }
     operands.push_back(codegen.redVal);
   }
   if (codegen.expValues)
     operands.push_back(codegen.expCount);
-  scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, lo, hi, step, operands);
+  scf::ForOp forOp = builder.create<scf::ForOp>(loc, lo, hi, step, operands);
   if (codegen.redVal)
     updateReduc(merger, codegen, forOp.getRegionIterArgs().front());
   if (codegen.expValues)
     codegen.expCount = forOp.getRegionIterArgs().back();
   // Assign induction variable to sparse or dense index.
   Value iv = forOp.getInductionVar();
   if (isSparse)
     codegen.pidxs[tensor][idx] = iv;
   else
     codegen.loops[idx] = iv;
-  rewriter.setInsertionPointToStart(forOp.getBody());
+  builder.setInsertionPointToStart(forOp.getBody());
   // Share vector iteration mask between all subsequent loads/stores.
   if (isVector)
-    codegen.curVecMask = genVectorMask(codegen, rewriter, iv, lo, hi, step);
+    codegen.curVecMask = genVectorMask(codegen, builder, iv, lo, hi, step);
   return forOp;
 }
 
 /// Emit a while-loop for co-iteration over multiple indices.
-static Operation *genWhile(Merger &merger, CodeGen &codegen,
-                           PatternRewriter &rewriter, linalg::GenericOp op,
-                           unsigned idx, bool needsUniv, BitVector &indices) {
+static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                           linalg::GenericOp op, unsigned idx, bool needsUniv,
+                           BitVector &indices) {
   SmallVector<Type, 4> types;
   SmallVector<Value, 4> operands;
   // Construct the while-loop with a parameter for each index.
-  Type indexType = rewriter.getIndexType();
+  Type indexType = builder.getIndexType();
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
     if (indices[b] && merger.isDim(b, Dim::kSparse)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
     }
   }
   if (codegen.redVal) {
     types.push_back(codegen.redVal.getType());
     operands.push_back(codegen.redVal);
   }
   if (codegen.expValues) {
     types.push_back(indexType);
     operands.push_back(codegen.expCount);
   }
   if (needsUniv) {
     types.push_back(indexType);
     operands.push_back(codegen.loops[idx]);
   }
   assert(types.size() == operands.size());
   Location loc = op.getLoc();
-  scf::WhileOp whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
+  scf::WhileOp whileOp = builder.create<scf::WhileOp>(loc, types, operands);
 
   SmallVector<Location> locs(types.size(), loc);
-  Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, types, locs);
-  Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, types, locs);
+  Block *before = builder.createBlock(&whileOp.getBefore(), {}, types, locs);
+  Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
-  rewriter.setInsertionPointToStart(&whileOp.getBefore().front());
+  builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
     if (indices[b] && merger.isDim(b, Dim::kSparse)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = before->getArgument(o);
       Value op2 = codegen.highs[tensor][idx];
-      Value opc = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
-                                                 op1, op2);
-      cond = cond ? rewriter.create<arith::AndIOp>(loc, cond, opc) : opc;
+      Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
+      cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
       codegen.pidxs[tensor][idx] = after->getArgument(o++);
     }
   }
   if (codegen.redVal)
     updateReduc(merger, codegen, after->getArgument(o++));
   if (codegen.expValues)
     codegen.expCount = after->getArgument(o++);
   if (needsUniv)
     codegen.loops[idx] = after->getArgument(o++);
   assert(o == operands.size());
-  rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
-  rewriter.setInsertionPointToStart(&whileOp.getAfter().front());
+  builder.create<scf::ConditionOp>(loc, cond, before->getArguments());
+  builder.setInsertionPointToStart(&whileOp.getAfter().front());
   return whileOp;
 }
 
 /// Generates a for-loop or a while-loop, depending on whether it implements
 /// singleton iteration or co-iteration over the given conjunction.
-static Operation *genLoop(Merger &merger, CodeGen &codegen,
-                          PatternRewriter &rewriter, linalg::GenericOp op,
-                          std::vector<unsigned> &topSort, unsigned at,
-                          bool needsUniv, BitVector &indices) {
+static Operation *genLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                          linalg::GenericOp op, std::vector<unsigned> &topSort,
+                          unsigned at, bool needsUniv, BitVector &indices) {
   unsigned idx = topSort[at];
   if (indices.count() == 1) {
     bool isOuter = at == 0;
     bool isInner = at == topSort.size() - 1;
-    return genFor(merger, codegen, rewriter, op, isOuter, isInner, idx,
-                  indices);
+    return genFor(merger, codegen, builder, op, isOuter, isInner, idx, indices);
   }
-  return genWhile(merger, codegen, rewriter, op, idx, needsUniv, indices);
+  return genWhile(merger, codegen, builder, op, idx, needsUniv, indices);
 }
 
 /// Generates the local variables for this loop, consisting of the sparse
 /// indices, restored universal dense index, and dense positions.
-static void genLocals(Merger &merger, CodeGen &codegen,
-                      PatternRewriter &rewriter, linalg::GenericOp op,
-                      std::vector<unsigned> &topSort, unsigned at,
-                      bool needsUniv, BitVector &locals) {
+static void genLocals(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                      linalg::GenericOp op, std::vector<unsigned> &topSort,
+                      unsigned at, bool needsUniv, BitVector &locals) {
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse indices.
   Value min;
   for (unsigned b = 0, be = locals.size(); b < be; b++) {
     if (locals[b] && merger.isDim(b, Dim::kSparse)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value ptr = codegen.indices[tensor][idx];
       Value s = codegen.pidxs[tensor][idx];
-      Value load = genLoad(codegen, rewriter, loc, ptr, s);
+      Value load = genLoad(codegen, builder, loc, ptr, s);
       codegen.idxs[tensor][idx] = load;
       if (!needsUniv) {
         if (min) {
-          Value cmp = rewriter.create<arith::CmpIOp>(
+          Value cmp = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ult, load, min);
-          min = rewriter.create<arith::SelectOp>(loc, cmp, load, min);
+          min = builder.create<arith::SelectOp>(loc, cmp, load, min);
         } else {
           min = load;
         }
       }
     }
   }
 
   // Merge dense universal index over minimum.
for (unsigned b = 0, be = locals.size(); b < be; b++) {
     if ((locals[b] || merger.isOutTensor(b, idx)) &&
         merger.isDim(b, Dim::kDense)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       unsigned pat = at;
       for (; pat != 0; pat--)
         if (codegen.pidxs[tensor][topSort[pat - 1]])
           break;
-      Value p = (pat == 0) ? constantIndex(rewriter, loc, 0)
+      Value p = (pat == 0) ? constantIndex(builder, loc, 0)
                            : codegen.pidxs[tensor][topSort[pat - 1]];
       codegen.pidxs[tensor][idx] = genAddress(
-          codegen, rewriter, loc, codegen.sizes[idx], p, codegen.loops[idx]);
+          codegen, builder, loc, codegen.sizes[idx], p, codegen.loops[idx]);
     }
   }
 
   // Move the insertion indices in lexicographic index order. During access
   // pattern expansion, we can skip setting the innermost dimension.
   if (codegen.sparseOut && !codegen.expValues) {
-    Value pos = constantIndex(rewriter, loc, at);
-    rewriter.create<memref::StoreOp>(loc, codegen.loops[idx], codegen.lexIdx,
-                                     pos);
+    Value pos = constantIndex(builder, loc, at);
+    builder.create<memref::StoreOp>(loc, codegen.loops[idx], codegen.lexIdx,
   }
 }
 
 /// Generates the induction structure for a while-loop.
 static void genWhileInduction(Merger &merger, CodeGen &codegen,
-                              PatternRewriter &rewriter, linalg::GenericOp op,
+                              OpBuilder &builder, linalg::GenericOp op,
                               unsigned idx, bool needsUniv,
                               BitVector &induction, scf::WhileOp whileOp) {
   Location loc = op.getLoc();
   // Finalize each else branch of all if statements.
   if (codegen.redVal || codegen.expValues) {
     while (auto ifOp = dyn_cast_or_null<scf::IfOp>(
-               rewriter.getInsertionBlock()->getParentOp())) {
+               builder.getInsertionBlock()->getParentOp())) {
       unsigned y = 0;
       SmallVector<Value, 4> yields;
       if (codegen.redVal) {
         yields.push_back(codegen.redVal);
         updateReduc(merger, codegen, ifOp.getResult(y++));
       }
       if (codegen.expValues) {
         yields.push_back(codegen.expCount);
         codegen.expCount = ifOp->getResult(y++);
       }
       assert(y == yields.size());
-      rewriter.create<scf::YieldOp>(loc, yields);
-      rewriter.setInsertionPointAfter(ifOp);
+      builder.create<scf::YieldOp>(loc, yields);
+      builder.setInsertionPointAfter(ifOp);
     }
   }
-  rewriter.setInsertionPointToEnd(&whileOp.getAfter().front());
+  builder.setInsertionPointToEnd(&whileOp.getAfter().front());
   // Finalize the induction. Note that the induction could be performed
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value, 4> operands;
-  Value one = constantIndex(rewriter, loc, 1);
+  Value one = constantIndex(builder, loc, 1);
   for (unsigned b = 0, be = induction.size(); b < be; b++) {
     if (induction[b] && merger.isDim(b, Dim::kSparse)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = codegen.idxs[tensor][idx];
       Value op2 = codegen.loops[idx];
       Value op3 = codegen.pidxs[tensor][idx];
-      Value cmp = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                                 op1, op2);
-      Value add = rewriter.create<arith::AddIOp>(loc, op3, one);
-      operands.push_back(rewriter.create<arith::SelectOp>(loc, cmp, add, op3));
+      Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+      Value add = builder.create<arith::AddIOp>(loc, op3, one);
+      operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, op3));
       codegen.pidxs[tensor][idx] = whileOp->getResult(o++);
     }
   }
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, whileOp->getResult(o++));
   }
   if (codegen.expValues) {
     operands.push_back(codegen.expCount);
     codegen.expCount = whileOp->getResult(o++);
   }
   if (needsUniv) {
     operands.push_back(
-        rewriter.create<arith::AddIOp>(loc, codegen.loops[idx], one));
+        builder.create<arith::AddIOp>(loc, codegen.loops[idx], one));
     codegen.loops[idx] = whileOp->getResult(o++);
   }
   assert(o == operands.size());
-  rewriter.create<scf::YieldOp>(loc, operands);
-  rewriter.setInsertionPointAfter(whileOp);
+  builder.create<scf::YieldOp>(loc, operands);
+  builder.setInsertionPointAfter(whileOp);
 }
 
 /// Generates the induction structure for a for-loop.
 static void genForInduction(Merger &merger, CodeGen &codegen,
-                            PatternRewriter &rewriter, linalg::GenericOp op,
+                            OpBuilder &builder, linalg::GenericOp op,
                             Operation *loop) {
   Location loc = op.getLoc();
   unsigned o = 0;
   SmallVector<Value, 4> operands;
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, loop->getResult(o++));
   }
   if (codegen.expValues) {
     operands.push_back(codegen.expCount);
     codegen.expCount = loop->getResult(o++);
   }
   assert(o == operands.size());
   if (o > 0)
-    rewriter.create<scf::YieldOp>(loc, operands);
-  rewriter.setInsertionPointAfter(loop);
+    builder.create<scf::YieldOp>(loc, operands);
+  builder.setInsertionPointAfter(loop);
 }
 
 /// Generates a single if-statement within a while-loop.
-static scf::IfOp genIf(Merger &merger, CodeGen &codegen,
-                       PatternRewriter &rewriter, linalg::GenericOp op,
-                       unsigned idx, BitVector &conditions) {
+static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                       linalg::GenericOp op, unsigned idx,
+                       BitVector &conditions) {
   Location loc = op.getLoc();
   SmallVector<Type, 4> types;
   Value cond;
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
     if (conditions[b]) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value clause;
       if (merger.isDim(b, Dim::kSparse)) {
         Value op1 = codegen.idxs[tensor][idx];
         Value op2 = codegen.loops[idx];
-        clause = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
-                                                op1, op2);
+        clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
       } else {
-        clause = constantI1(rewriter, loc, true);
+        clause = constantI1(builder, loc, true);
       }
-      cond = cond ? rewriter.create<arith::AndIOp>(loc, cond, clause) : clause;
+      cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
     }
   }
   if (codegen.redVal)
     types.push_back(codegen.redVal.getType());
   if (codegen.expValues)
-    types.push_back(rewriter.getIndexType());
-  scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, types, cond, /*else=*/true);
-  rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
+    types.push_back(builder.getIndexType());
+  scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, cond, /*else=*/true);
+  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   return ifOp;
 }
 
 /// Generates end of true branch of if-statement within a while-loop.
-static void endIf(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
+static void endIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                   linalg::GenericOp op, scf::IfOp ifOp, Operation *loop,
                   Value redInput, Value cntInput) {
   SmallVector<Value, 4> operands;
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, redInput);
   }
   if (codegen.expValues) {
     operands.push_back(codegen.expCount);
     codegen.expCount = cntInput;
   }
   if (!operands.empty())
-    rewriter.create<scf::YieldOp>(op.getLoc(), operands);
-  rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+    builder.create<scf::YieldOp>(op.getLoc(), operands);
+  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (loop sequence).
 //===----------------------------------------------------------------------===//
 
 /// Starts a loop sequence at given level. Returns true if
 /// the universal loop index must be maintained at this level.
-static bool startLoopSeq(Merger &merger, CodeGen &codegen,
-                         PatternRewriter &rewriter, linalg::GenericOp op,
-                         std::vector<unsigned> &topSort, unsigned exp,
-                         unsigned at, unsigned idx, unsigned ldx,
+static bool startLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                         linalg::GenericOp op, std::vector<unsigned> &topSort,
+                         unsigned exp, unsigned at, unsigned idx, unsigned ldx,
                          unsigned lts) {
   assert(codegen.curVecLength == 1);
   assert(!codegen.loops[idx]);
   // Emit invariants at this loop sequence level.
-  genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/true);
+  genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/true);
   // Emit access pattern expansion for sparse tensor output.
-  genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/true);
+  genExpansion(merger, codegen, builder, op, at, /*atStart=*/true);
   // Emit further intitialization at this loop sequence level.
   unsigned l0 = merger.set(lts)[0];
   bool needsUniv =
-      genInit(merger, codegen, rewriter, op, topSort, at, merger.lat(l0).bits);
+      genInit(merger, codegen, builder, op, topSort, at, merger.lat(l0).bits);
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = merger.set(lts).size();
     for (unsigned i = 1; i < lsize; i++) {
       unsigned li = merger.set(lts)[i];
       if (!merger.hasAnyDimOf(merger.lat(li).simple, Dim::kSparse))
         return true;
     }
   }
   return false;
 }
 
 /// Starts a single loop in current sequence.
 static Operation *startLoop(Merger &merger, CodeGen &codegen,
-                            PatternRewriter &rewriter, linalg::GenericOp op,
+                            OpBuilder &builder, linalg::GenericOp op,
                             std::vector<unsigned> &topSort, unsigned at,
                             unsigned li, bool needsUniv) {
   assert(codegen.curVecLength == 1);
   // Emit the for/while-loop control.
-  Operation *loop = genLoop(merger, codegen, rewriter, op, topSort, at,
+  Operation *loop = genLoop(merger, codegen, builder, op, topSort, at,
                             needsUniv, merger.lat(li).simple);
   // Emit the locals for this loop.
-  genLocals(merger, codegen, rewriter, op, topSort, at, needsUniv,
+  genLocals(merger, codegen, builder, op, topSort, at, needsUniv,
             merger.lat(li).bits);
   return loop;
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
-static bool endLoop(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
+static bool endLoop(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                     linalg::GenericOp op, Operation *loop, unsigned idx,
                     unsigned li, bool needsUniv) {
   codegen.curVecLength = 1;
   // End a while-loop.
   if (auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
-    genWhileInduction(merger, codegen, rewriter, op, idx, needsUniv,
+    genWhileInduction(merger, codegen, builder, op, idx, needsUniv,
                       merger.lat(li).bits, whileOp);
     return needsUniv;
   }
   // End a for-loop.
-  genForInduction(merger, codegen, rewriter, op, loop);
+  genForInduction(merger, codegen, builder, op, loop);
   return false;
 }
 
 /// Ends a loop sequence at given level.
-static void endLoopSeq(Merger &merger, CodeGen &codegen,
-                       PatternRewriter &rewriter, linalg::GenericOp op,
-                       unsigned exp, unsigned at, unsigned idx, unsigned ldx) {
+static void endLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
+                       linalg::GenericOp op, unsigned exp, unsigned at,
+                       unsigned idx, unsigned ldx) {
   assert(codegen.curVecLength == 1);
   codegen.loops[idx] = Value();
   // Bring a pending reduction back from SIMD form when sequence ends.
   if (codegen.redVal)
     if (auto vtp = codegen.redVal.getType().dyn_cast<VectorType>())
       updateReduc(merger, codegen,
-                  genVectorReducEnd(codegen, rewriter, op.getLoc(), vtp));
+                  genVectorReducEnd(codegen, builder, op.getLoc(), vtp));
   // Unmark bookkeeping of invariants and loop index.
-  genInvariants(merger, codegen, rewriter, op, exp, ldx, /*atStart=*/false);
+  genInvariants(merger, codegen, builder, op, exp, ldx, /*atStart=*/false);
   // Finalize access pattern expansion for sparse tensor output.
-  genExpansion(merger, codegen, rewriter, op, at, /*atStart=*/false);
+  genExpansion(merger, codegen, builder, op, at, /*atStart=*/false);
 }
 
 /// Recursively generates code while computing iteration lattices in order
 /// to manage the complexity of implementing co-iteration over unions
 /// and intersections of sparse iterations spaces.
-static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
+static void genStmt(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
                     linalg::GenericOp op, std::vector<unsigned> &topSort,
                     unsigned exp, unsigned at) {
   // At each leaf, assign remaining tensor (sub)expression to output tensor.
   if (at == topSort.size()) {
     unsigned ldx = topSort[at - 1];
     Value rhs = genExp(merger, codegen, rewriter, op, exp, ldx);
     genTensorStore(merger, codegen, rewriter, op, exp, rhs);
     return;
needsUniv =
         endLoop(merger, codegen, rewriter, op, loop, idx, li, needsUniv);
   }
 
   // End a loop sequence.
   endLoopSeq(merger, codegen, rewriter, op, exp, at, idx, ldx);
 }
 
 /// Converts the result computed by the sparse kernel into the required form.
-static void genResult(Merger &merger, CodeGen &codegen,
-                      PatternRewriter &rewriter, linalg::GenericOp op) {
+static void genResult(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
+                      linalg::GenericOp op) {
   OpOperand *lhs = op.getOutputOperand(0);
   Type resType = lhs->get().getType();
   if (getSparseTensorEncoding(resType)) {
     // The sparse tensor rematerializes from the original sparse tensor's
     // underlying sparse storage format.
     rewriter.replaceOpWithNewOp<LoadOp>(op, resType, lhs->get(),
                                         codegen.sparseOut == lhs);
   } else {

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

if (x.hasValue() && y.hasValue()) {
       if (isa<sparse_tensor::BinaryOp>(def))
         return addExp(kBinary, e0, e1, Value(), def);
     }
   }
   // Cannot build.
   return None;
 }
 
-static Value insertYieldOp(PatternRewriter &rewriter, Location loc,
-                           Region &region, ValueRange vals) {
+static Value insertYieldOp(RewriterBase &rewriter, Location loc, Region &region,
+                           ValueRange vals) {
   // Make a clone of overlap region.
   Region tmpRegion;
   BlockAndValueMapping mapper;
   region.cloneInto(&tmpRegion, tmpRegion.begin(), mapper);
   Block &clonedBlock = tmpRegion.front();
   YieldOp clonedYield = cast<YieldOp>(clonedBlock.getTerminator());
   // Merge cloned block and return yield value.
   Operation *placeholder = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   rewriter.mergeBlockBefore(&tmpRegion.front(), placeholder, vals);
   Value val = clonedYield.result();
   rewriter.eraseOp(clonedYield);
   rewriter.eraseOp(placeholder);
   return val;
 }
 
-static Value buildUnaryPresent(PatternRewriter &rewriter, Location loc,
+static Value buildUnaryPresent(RewriterBase &rewriter, Location loc,
                                Operation *op, Value v0) {
   if (!v0)
     // Empty input value must be propagated.
     return Value();
   UnaryOp unop = cast<UnaryOp>(op);
   Region &presentRegion = unop.presentRegion();
   if (presentRegion.empty())
     // Uninitialized Value() will be interpreted as missing data in the
     // output.
     return Value();
   return insertYieldOp(rewriter, loc, presentRegion, {v0});
 }
 
-static Value buildBinaryOverlap(PatternRewriter &rewriter, Location loc,
+static Value buildBinaryOverlap(RewriterBase &rewriter, Location loc,
                                 Operation *op, Value v0, Value v1) {
   if (!v0 || !v1)
     // Empty input values must be propagated.
     return Value();
   BinaryOp binop = cast<BinaryOp>(op);
   Region &overlapRegion = binop.overlapRegion();
   if (overlapRegion.empty())
     // Uninitialized Value() will be interpreted as missing data in the
     // output.
     return Value();
   return insertYieldOp(rewriter, loc, overlapRegion, {v0, v1});
 }
 
-Value Merger::buildExp(PatternRewriter &rewriter, Location loc, unsigned e,
+Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
                        Value v0, Value v1) {
   switch (tensorExps[e].kind) {
   case kTensor:
   case kInvariant:
   case kIndex:
     llvm_unreachable("unexpected non-op");
   // Unary ops.
   case kAbsF: