[mlir][sparse] use macro to build common operations in LoopEmitter

[aartbik]

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

40–42

I am not sure if we have a style recommendation for this, but I use it myself at lots of places since I like how short it gets.
I would just a comment here

// A few shorthand macros specific to this file.

or so

[Peiming]

address comments.

[Peiming]

more macros.

[Peiming]

add macro for yield op

[Peiming]

rebase

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

 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
-// File local helper functions.
+// File local shorthand macros
 //===----------------------------------------------------------------------===//
 
 #define CMPI(p, l, r)                                                          \
-  (builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::p, l, r)           \
+  (builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::p, (l), (r))       \
        .getResult())
 
-#define ADDI(lhs, rhs) (builder.create<arith::AddIOp>(loc, lhs, rhs))
+#define C_IDX(v) (constantIndex(builder, loc, (v)))
+#define YIELD(vs) (builder.create<scf::YieldOp>(loc, (vs)))
+#define ADDI(lhs, rhs) (builder.create<arith::AddIOp>(loc, (lhs), (rhs)))
+#define ANDI(lhs, rhs) (builder.create<arith::AndIOp>(loc, (lhs), (rhs)))
+#define SUBI(lhs, rhs) (builder.create<arith::SubIOp>(loc, (lhs), (rhs)))
+#define MULI(lhs, rhs) (builder.create<arith::MulIOp>(loc, (lhs), (rhs)))
+#define SELECT(c, l, r) (builder.create<arith::SelectOp>(loc, (c), (l), (r)))
 
-#define C_IDX(v) (constantIndex(builder, loc, v))
+//===----------------------------------------------------------------------===//
+// File local helper functions.
+//===----------------------------------------------------------------------===//
 
 /// 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 genIndexLoad(OpBuilder &builder, Location loc, Value mem,
                           Value s) {
   // For the scalar case, we simply zero extend narrower indices into 64-bit
   // values before casting to index without a performance penalty. Here too,

 
 /// Converts a coordinate relative to the slice to the coordinate relative
 /// to the underlying tensor.
 // FIXME: that description says "sliceCrd -> tensorCrd"; but the function
 // name suggests it should be "tensorCrd -> sliceCrd".
 static Value toSliceCrd(OpBuilder &builder, Location loc, Value crd,
                         Value offset, Value stride, Value tensor, Level lvl) {
   // tensorCrd = sliceCrd * stride + offset
-  crd = builder.create<arith::MulIOp>(loc, crd, stride);
-  crd = builder.create<arith::AddIOp>(loc, crd, offset);
-  return crd;
+  return ADDI(MULI(crd, stride), offset);
 }
 
 /// Generates code to compute the *absolute* offset of the slice based on the
 /// provide minimum coordinates in the slice.
 /// E.g., when reducing d0 + d1 + d2, we need two slices to fully reduced the
 /// expression, i,e, s1 = slice(T, d0), s2 = slice(s1, d1). The *absolute*
 /// offset is the offset computed relative to the initial tensors T.
 ///
 /// When isNonEmpty == true, the computed offset is meaningless and should not
 /// be used during runtime, the method generates code to return 0 currently in
 /// that case.
 ///
 /// offset = isNonEmpty && minCrd >= size ? minCrd - size + 1 : 0;
 static Value offsetFromMinCoord(OpBuilder &builder, Location loc, Value minCrd,
                                 Value size, Value isNonEmpty) {
   Value geSize = CMPI(uge, minCrd, size);
-  Value pred = builder.create<arith::AndIOp>(loc, isNonEmpty, geSize);
-  Value mp1 = builder.create<arith::AddIOp>(loc, minCrd, C_IDX(1));
-  Value mms = builder.create<arith::SubIOp>(loc, mp1, size);
+  Value pred = ANDI(isNonEmpty, geSize);
+  // Computes minCrd - size + 1
+  Value mms = SUBI(ADDI(minCrd, C_IDX(1)), size);
   // This is the absolute offset related to the underly tensor.
-  return builder.create<arith::SelectOp>(loc, pred, mms, C_IDX(0));
+  return SELECT(pred, mms, C_IDX(0));
 }
 
 /// Converts a coordinate relative to the underlying tensor to the coordinate
 /// relative to the slice, returns a extra reminder value
 // FIXME: that description says "tensorCrd -> sliceCrd"; but the function
 // name suggests it should be "sliceCrd -> tensorCrd".
 static std::pair<Value, Value> fromSliceCrd(OpBuilder &builder, Location loc,
                                             Value crd, Value offset,
                                             Value stride, Value tensor,
                                             Level lvl) {
   // sliceCrd = (tensorCrd - offset) / stride
-  crd = builder.create<arith::SubIOp>(loc, crd, offset);
+  crd = SUBI(crd, offset);
   Value rem = builder.create<arith::RemUIOp>(loc, crd, stride);
   crd = builder.create<arith::DivUIOp>(loc, crd, stride);
   return std::make_pair(crd, rem);
 }
 
 std::pair<Value, Value>
 LoopEmitter::genSliceLegitPredicate(OpBuilder &builder, Location loc, Value crd,
                                     TensorId tid, Level lvl) {
if (auto staticStride = enc.getStaticLvlSliceStride(lvl);
       !(staticStride.has_value() && *staticStride == 1)) {
     auto fitStride = CMPI(eq, crdRem, C_IDX(0));
     conds.push_back(fitStride);
   }
 
   // Must meet all condition to be a valid coordinate in slice.
   auto pred = conds.front();
   for (auto cond : ValueRange(conds).drop_front())
-    pred = builder.create<arith::AndIOp>(loc, pred, cond);
+    pred = ANDI(pred, cond);
 
   return {newCrd, pred};
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 Value LoopEmitter::genAddress(OpBuilder &builder, Location loc, TensorId tid,
                               Level lvl, Value crd) {
   Value pos = lvl == 0 ? C_IDX(0) : posits[tid][lvl - 1];
-  Value mul = builder.create<arith::MulIOp>(loc, highs[tid][lvl], pos);
+  Value mul = MULI(highs[tid][lvl], pos);
   if (isSparseSlices[tid])
     crd = toSliceCrd(builder, loc, crd, sliceOffsets[tid][lvl],
                      sliceStrides[tid][lvl], tensors[tid], lvl);
-  Value add = builder.create<arith::AddIOp>(loc, mul, crd);
+  Value add = ADDI(mul, crd);
   return add;
 }
 
 Value LoopEmitter::genSegmentHigh(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl, Value pLo,
                                   Value pHi) {
   const auto coordinates = coordinatesBuffers[tid][lvl];
   const auto sameCrd = genIndexLoad(builder, loc, coordinates, pLo);
auto whileOp = builder.create<scf::WhileOp>(
         {
           OpBuilder::InsertionGuard guard(builder);
           // Load the next coordinates only when inbound (to avoid OOB
           // acccesses).
           builder.setInsertionPointToStart(ifInBound.thenBlock());
           Value crd = genIndexLoad(builder, loc, coordinates, pos);
           Value isSameCrd = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::eq, crd, sameCrd);
-          builder.create<scf::YieldOp>(loc, isSameCrd);
+          YIELD(isSameCrd);
           // Else, the position is out of bound, yield false to terminate the
           // loop.
           builder.setInsertionPointToStart(ifInBound.elseBlock());
-          builder.create<scf::YieldOp>(loc, constantI1(builder, loc, false));
+          YIELD(constantI1(builder, loc, false));
         }
         builder.create<scf::ConditionOp>(loc, ifInBound.getResults()[0], ivs);
       },
       /*afterBuilder=*/
       [](OpBuilder &builder, Location loc, ValueRange ivs) {
         // pos ++
-        Value nextPos = builder.create<arith::AddIOp>(loc, ivs[0], C_IDX(1));
-        builder.create<scf::YieldOp>(loc, nextPos);
+        Value nextPos = ADDI(ivs[0], C_IDX(1));
+        YIELD(nextPos);
       });
   // Return the segment high.
   return whileOp.getResult(0);
 }
 
 Value LoopEmitter::genSparseCrd(OpBuilder &builder, Location loc, TensorId tid,
                                 Level dstLvl) {
   Value crd = C_IDX(0);
   const auto reassoc = getCollapseReassociation(tid, dstLvl);
   const unsigned reassocSize = reassoc.size();
   for (unsigned i = 0; i < reassocSize; i++) {
     const Level srcLvl = reassoc[i];
     // A load on the coordinates array yields the coordinate.
     const Value mem = coordinatesBuffers[tid][srcLvl];
     /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
     const Value pos = posits[tid][dstLvl];
     const Value off = genIndexLoad(builder, loc, mem, pos);
     // Linearized the coordinates within the same collapse reassociation.
-    crd = builder.create<arith::AddIOp>(loc, crd, off);
+    crd = ADDI(crd, off);
     if (i != reassocSize - 1) {
-      crd = builder.create<arith::MulIOp>(loc, crd,
-                                          this->lvlSizes[tid][reassoc[i + 1]]);
+      crd = MULI(crd, this->lvlSizes[tid][reassoc[i + 1]]);
     }
   }
   return crd;
 }
 
 LoopEmitter::LoopEmitter(ValueRange tensors, StringAttr loopTag, bool hasOutput,
                          bool isSparseOut, ArrayRef<LoopId> topSort,
                          DependentLvlGetter dimGetter) {
for (Level lvl = 0; lvl < lvlRank; lvl++) {
       if (!dependentLvlMap[t][lvl].empty()) {
         ArrayRef<std::pair<TensorId, Level>> depLvls = dependentLvlMap[t][lvl];
         // Needs at least two operands to form a non-trivial affine expression.
         assert(depLvls.size() > 1);
 
         Value size = c0;
         for (unsigned e = depLvls.size() - 1; e >= 1; e--) {
           auto [dt, dd] = depLvls[e];
-          size = builder.create<arith::AddIOp>(loc, size, lvlSizes[dt][dd]);
+          size = ADDI(size, lvlSizes[dt][dd]);
           sliceSizes[t][lvl][e - 1] = size;
         }
       }
     }
   }
   localInsertPos = builder.getInsertionPoint()->getPrevNode();
 }
 
for (auto [tid, lvl, res] : slicedTids) {
     } else {
       if (!isDenseDLT(lvlTypes[tid][lvl])) {
         // Else this is a resolved-slice, and advance posit similar to TACO.
         Value c1 = C_IDX(1), c2 = C_IDX(2);
         // pIdx += 2, we finished the current lvl, advance the pointer index of
         // the previous level by two to skip the [pLo, pHi] for current level.
         Value sPtrBuf = slicePosBuffer[tid][lvl].back();
         Value curP = genIndexLoad(builder, loc, sPtrBuf, c1);
-        Value nexP = builder.create<arith::AddIOp>(loc, curP, c2);
         // TODO: we could probably use an SSA value for it.
+        Value nexP = ADDI(curP, c2);
         builder.create<memref::StoreOp>(loc, nexP, sPtrBuf, c1);
       }
     }
   }
   loopSeqStack.pop_back();
 }
 
 Value LoopEmitter::genAffine(OpBuilder &builder, Location loc, AffineExpr a) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     // FIXME: since the one callsite in Sparsification passes in a
     // level-expression, the `getPosition` must in fact be a `Dimension`.
     // However, elsewhere we have been lead to expect that `loopIdToOrd`
     // should be indexed by `LoopId`...
     const auto loopId = a.cast<AffineDimExpr>().getPosition();
     assert(loopId < loopIdToOrd.size());
     return loopStack[loopIdToOrd[loopId]].iv;
   }
   case AffineExprKind::Add: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return builder.create<arith::AddIOp>(
-        loc, genAffine(builder, loc, binOp.getLHS()),
+    return ADDI(genAffine(builder, loc, binOp.getLHS()),
                 genAffine(builder, loc, binOp.getRHS()));
   }
   case AffineExprKind::Mul: {
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return builder.create<arith::MulIOp>(
-        loc, genAffine(builder, loc, binOp.getLHS()),
+    return MULI(genAffine(builder, loc, binOp.getLHS()),
                 genAffine(builder, loc, binOp.getRHS()));
   }
   case AffineExprKind::Constant: {
     int64_t c = a.cast<AffineConstantExpr>().getValue();
     return C_IDX(c);
   }
   default:
     llvm_unreachable("unexpected affine subscript");
   }
scf::IfOp ifOp = builder.create<scf::IfOp>(loc, types, pred,
                                                /*else*/ hasReduc);
     if (hasReduc) {
       // scf.for (a) -> v
       //  %s = scf.if (a) -> v
       //    user-generated code.
       //  else
       //    yield a
       //  yield %s
-      builder.create<scf::YieldOp>(loc, ifOp.getResults());
+      YIELD(ifOp.getResults());
       builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
       // On mismatch.
-      builder.create<scf::YieldOp>(loc, reduc);
+      YIELD(reduc);
     }
     // Set the insertion point to matched branch.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
     crd = trans;
   }
 
   assert(crd);
   coords[tid][dstLvl] = crd;
if (isDenseSliceCond) {
       bool fullyReduced = depFullyReduced(tid, lvl);
       Value sliceSz = sliceSizes[tid][lvl][sliceStack[tid].back().depth - 1];
       // Adjust for loop hi for dense slice-driven loop.
       if (fullyReduced) {
         hi = sliceSz;
         condTid.push_back(tid);
         condLvl.push_back(lvl);
       } else {
-        hi = builder.create<arith::SubIOp>(loc, lvlSizes[tid][lvl], sliceSz);
-        hi = builder.create<arith::AddIOp>(loc, hi, C_IDX(1));
+        hi = SUBI(lvlSizes[tid][lvl], sliceSz);
+        hi = ADDI(hi, C_IDX(1));
       }
     } else {
       condTid.push_back(tid);
       condLvl.push_back(lvl);
     }
     l = emitForLoopOverTensorAtLvl(builder, loc, tid, lvl, lo, hi, reduc,
                                    isParallel);
   }
   Value iv = coords[tid][lvl];
   for (auto [t, l] : llvm::zip(tids, lvls)) {
     // We only need to handle slice-driven loops on dense level here.
     // If it is a slice-driven loop on sparse level, it needs a while loop to
     // insert break statements, and it must have been handled correctly in L692.
     if (!dependentLvlMap[t][l].empty() && isDenseDLT(lvlTypes[t][l])) {
       // Pushes sliced levels to build correct LoopInfo.
       bool fullyReduc = depFullyReduced(t, l);
       SliceInfo &info = sliceStack[t].back();
       if (fullyReduc) {
-        posits[t][l] =
-            genAddress(builder, loc, t, l,
-                       builder.create<arith::AddIOp>(loc, info.offset, iv));
+        posits[t][l] = genAddress(builder, loc, t, l, ADDI(info.offset, iv));
       } else {
         // Puts sliced dense loop into LoopInfo so that LoopEmitter knows how to
         // exit it.
         sliceTids.push_back(t);
         sliceLvls.push_back(l);
         sliceReduc.push_back(fullyReduc);
         // Update the slice information as we enter the new loop.
         assert(*info.slicedOnLvl == l);
scf::IfOp ifOp =
       builder.create<scf::IfOp>(loc, types, pred, /*else*/ hasReduc);
   if (hasReduc) {
     // scf.for (a) -> v
     //  %s = scf.if (a) -> v
     //    user-generated code.
     //  else
     //    yield a
     //  yield %s
-    builder.create<scf::YieldOp>(loc, ifOp.getResults());
+    YIELD(ifOp.getResults());
     builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
     // On mismatch.
-    builder.create<scf::YieldOp>(loc, reduc);
+    YIELD(reduc);
   }
   // Set the insert point to matched branch.
   builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
 
   // NOTE: we can also prepare for next lvl here in advance
   // Push the loop into stack
   loopStack.emplace_back(ArrayRef<TensorId>(tid), ArrayRef<Level>(lvl),
                          ArrayRef<TensorId>(), ArrayRef<Level>(),
if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
           // Links the SSA chain for segHi.
           segHi[tid][reassoc[i]] = after->getArgument(o++);
         }
       }
       Value op1 = before->getArgument(o);
       // We used the first level bound as the bound the collapsed set of levels.
       Value op2 = highs[tid][reassoc.front()];
       Value opc = CMPI(ult, op1, op2);
-      cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
+      cond = cond ? ANDI(cond, opc) : opc;
       // Update positions
       Value pos = after->getArgument(o++);
       // For COO, the position is the same across consecutive levels.
       /// FIXME: See the [CLARIFY_POSITS_LVL] note in the header.
       llvm::for_each(reassoc, [this, tid, pos](Level srcLvl) {
         posits[tid][srcLvl] = pos;
       });
     }
if (!slicesPreds.empty()) {
     SmallVector<Value> yields(after->getArguments());
     // Generates a list of if statments
     //  pos = in_slice ? pos : pos + 1
     // TODO: instead of always picking pos + 1, we should set pos = high to
     // break to loop if the coordinates are larger than the slice size.
     //
     // This "idx" is the index into `llvm::zip(tids, lvls)`
     for (auto [pred, idx] : slicesPreds) {
-      Value nextPos = builder.create<arith::AddIOp>(loc, yields[idx], C_IDX(1));
-      yields[idx] =
-          builder.create<arith::SelectOp>(loc, pred, yields[idx], nextPos);
+      Value nextPos = ADDI(yields[idx], C_IDX(1));
+      yields[idx] = SELECT(pred, yields[idx], nextPos);
     }
 
     Value pred = slicesPreds.front().first;
     for (int i = 1, e = slicesPreds.size(); i < e; i++) {
-      pred = builder.create<arith::AndIOp>(loc, pred, slicesPreds[i].first);
+      pred = ANDI(pred, slicesPreds[i].first);
     }
     auto ifOp = builder.create<scf::IfOp>(loc, types, pred, /*else*/ true);
     ifOp->setAttr(getLoopEmitterLoopAttrName(),
                   StringAttr::get(builder.getContext(), "slice"));
-    builder.create<scf::YieldOp>(loc, ifOp->getResults());
+    YIELD(ifOp->getResults());
     assert(types.size() == yields.size());
     // If not all slices are legit
     builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
-    builder.create<scf::YieldOp>(loc, yields);
+    YIELD(yields);
 
     // If all slices are legit, start the user generated code.
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
   }
 
   Value min;
   // Finds the minimum coordinate
   if (!needsUniv) {
     for (auto [tid, lvl] : llvm::zip(tids, lvls)) {
       const auto lvlTp = lvlTypes[tid][lvl];
       if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
         const auto crd = coords[tid][lvl];
         if (min) {
-          Value cmp = CMPI(ult, crd, min);
-          min = builder.create<arith::SelectOp>(loc, cmp, crd, min);
+          Value cmp = CMPI(ult, coords[tid][lvl], min);
+          min = SELECT(cmp, coords[tid][lvl], min);
         } else {
           min = crd;
         }
       }
     }
   } else {
     assert(!min);
     // Otherwise, universal index is the minimal pos.
for (unsigned i = 0, e = reassoc.size() - 1; i < e; i++) {
         assert(oldSegHi);
         Value newSegHi = builder.create<arith::CmpIOp>(
             loc, arith::CmpIPredicate::uge, pos, oldSegHi);
         auto ifNewSegHi = builder.create<scf::IfOp>(loc, builder.getIndexType(),
                                                     newSegHi, true);
         {
           OpBuilder::InsertionGuard guard(builder);
           builder.setInsertionPointToStart(ifNewSegHi.thenBlock());
-          builder.create<scf::YieldOp>(loc,
-                                       genSegmentHigh(builder, loc, tid, srcLvl,
-                                                      pos, highs[tid][srcLvl]));
+          YIELD(genSegmentHigh(builder, loc, tid, srcLvl, pos,
+                               highs[tid][srcLvl]));
           // Else, resues the same segment high.
           builder.setInsertionPointToStart(ifNewSegHi.elseBlock());
-          builder.create<scf::YieldOp>(loc, oldSegHi);
+          YIELD(oldSegHi);
         }
         highs[tid][srcLvl + 1] = segHi[tid][srcLvl] = ifNewSegHi.getResult(0);
       }
     };
     const auto srcLvl = reassoc.back();
     if (!isUniqueDLT(lvlTypes[tid][srcLvl])) {
       segHi[tid][srcLvl] = genSegmentHigh(
           builder, loc, tid, srcLvl, posits[tid][srcLvl], highs[tid][srcLvl]);
for (const Level srcLvl : getCollapseReassociation(tid, dstLvl)) {
     if (!isCompressedDLT(lvlTp) && !isSingletonDLT(lvlTp))
       continue;
     if (isCompressedDLT(lvlTp)) {
       const Value mem = positionsBuffers[tid][srcLvl];
 
       const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
       posits[tid][srcLvl] = genIndexLoad(builder, loc, mem, pLo);
 
-      const Value pHi = builder.create<arith::AddIOp>(loc, pLo, c1);
+      const Value pHi = ADDI(pLo, c1);
       highs[tid][srcLvl] = genIndexLoad(builder, loc, mem, pHi);
       return;
     }
     if (isSingletonDLT(lvlTp)) {
       const Value pLo = srcLvl == 0 ? c0 : posits[tid][srcLvl - 1];
       posits[tid][srcLvl] = pLo;
 
       // If we are coiterating non-unique levels, then use pHi=segHi;
       // otherwise use pHi=pLo+1.
       // NOTE: Just because the level is non-unique, that does not
       // guarantee that segHi is defined: because we only generate segHi
       // whenever coiterating, in order to improve code quality for the
       // non-coiterating cases.
       const auto parentSegHi = segHi[tid][srcLvl - 1];
       highs[tid][srcLvl] =
           (!isUniqueDLT(lvlTypes[tid][srcLvl - 1]) && parentSegHi)
               ? parentSegHi
-              : builder.create<arith::AddIOp>(loc, pLo, c1);
+              : ADDI(pLo, c1);
       return;
     }
   }
 
   llvm_unreachable("Unrecognized level-type!");
 }
 
 void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(OpBuilder &builder,
if (isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp)) {
       const Value crd = coords[tid][dstLvl];
       const Value pos = posits[tid][dstLvl];
       Value cmp = CMPI(eq, crd, iv);
       // If the loop contains a coiteration with non-unique level, we fast
       // forward all the duplicated coords by setting the position to the
       // segment high.
       Value add = !isUniqueDLT(lvlTypes[tid][reassoc.back()])
                       ? segHi[tid][reassoc.back()]
-                      : builder.create<arith::AddIOp>(loc, pos, one);
+                      : ADDI(pos, one);
 
-      operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, pos));
+      operands.push_back(SELECT(cmp, add, pos));
       // Following loops continue iteration from the break point of the
       // current while loop.
       const Value newPos = whileOp->getResult(o++);
       // We need to define a new local variable for `tid` to avoid
       // warnings about "captured structured bindings are a C++20 extension".
       // FIXME(wrengr): define a helper function to capture this idiom!
       const TensorId newTid = tid;
       llvm::for_each(reassoc, [this, newTid, newPos](Level srcLvl) {
for (auto &i : reduc) {
     // In place update reduction variable.
     i = whileOp->getResult(o++);
   }
 
   // An (optional) universal index.
   if (operands.size() + delta < whileOp.getNumResults()) {
     assert(operands.size() + delta + 1 == whileOp.getNumResults());
     // The last one is the universial index.
-    operands.push_back(builder.create<arith::AddIOp>(loc, iv, one));
+    operands.push_back(ADDI(iv, one));
     // update the loop starting point of current loop sequence
     loopSeqStack.back().first = whileOp->getResult(o++);
   }
 
   assert(o == operands.size() + delta);
-  builder.create<scf::YieldOp>(loc, operands);
+  YIELD(operands);
   builder.setInsertionPointAfter(whileOp);
 }
 
 void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,
                                   MutableArrayRef<Value> reduc) {
   // Clean up the values, it would help use to discover potential bug at a
   // earlier stage (instead of silently using a wrong value).
   const LoopInfo &loopInfo = loopStack.back();

 //   }
 //   loopLo ++;
 // }
 std::pair<Operation *, ValueRange> LoopEmitter::genSliceLvlTraverseLoop(
     OpBuilder &builder, Location loc, Value loopLo, Value loopHi, Value offset,
     Value size, TensorId tid, Level lvl, ValueRange userReduc, bool genYield,
     LoopBodyBuilder bodyBuilder) {
   Value c1 = C_IDX(1);
-  Value sliceHi = builder.create<arith::AddIOp>(loc, offset, size);
+  Value sliceHi = ADDI(offset, sliceSizes[tid][lvl].back());
 
   SmallVector<Value> reduc = {
       loopLo,                         // loop lower bounds
       constantI1(builder, loc, true), // continue
   };
   // Append user required reduction value.
   reduc.append(userReduc.begin(), userReduc.end());
   scf::WhileOp whileOp = builder.create<scf::WhileOp>(
       loc, ValueRange(reduc).getTypes(), reduc,
       /*beforeBuilder=*/
       [loopHi](OpBuilder &builder, Location loc, ValueRange args) {
         Value lo = args[0];
         Value cont = args[1];
         Value inBound = CMPI(ult, lo, loopHi);
-        Value cond = builder.create<arith::AndIOp>(loc, cont, inBound);
+        Value cond = ANDI(cont, inBound);
         // continue if not yet break nor out of bound.
         builder.create<scf::ConditionOp>(loc, cond, args);
       },
       /*afterBuilder=*/
       [this, c1, tid, lvl, sliceHi, genYield,
        bodyBuilder](OpBuilder &builder, Location loc, ValueRange args) {
         Value iv = args[0];
         Value coord =
             genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], iv);
         Value cont = CMPI(ult, coord, sliceHi);
         TypeRange types = args.drop_front(2).getTypes();
 
         auto ifOp = builder.create<scf::IfOp>(loc, types, cont, true);
         {
           // 2 reduction variable maintained by us.
           SmallVector<Value> ifRet = args.drop_front(2);
           assert(ifRet.size() == args.size() - 2);
 
           OpBuilder::InsertionGuard guard(builder);
           // If coord >= sliceHi.
           builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
-          // Coordinates is OOB, just yield.
-          builder.create<scf::YieldOp>(loc, ifRet);
+          YIELD(ifRet);
 
           // If coord < sliceHi.
           builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
           // Delegates to users' callback.
           bodyBuilder(builder, loc, iv, ifRet);
           if (genYield) {
             builder.setInsertionPointToEnd(&ifOp.getThenRegion().front());
-            builder.create<scf::YieldOp>(loc, ifRet);
+            YIELD(ifRet);
           }
         }
         // Marks this speical ifOp to avoid sparisification finalizing it.
         ifOp->setAttr(getLoopEmitterLoopAttrName(),
                       StringAttr::get(builder.getContext(), "slice"));
         // Insertion point restored to after ifOp.
         SmallVector<Value> yields;
         // Increase induction variable.
-        yields.push_back(builder.create<arith::AddIOp>(loc, iv, c1));
-        // Terminates the while loop according to the continue flag.
+        yields.push_back(ADDI(iv, c1));
         yields.push_back(cont);
         yields.append(ifOp.getResults().begin(), ifOp.getResults().end());
-        builder.create<scf::YieldOp>(loc, yields);
+        YIELD(yields);
       });
 
   builder.setInsertionPointAfter(whileOp);
   return std::make_pair(whileOp, whileOp.getResults().drop_front(2));
 }
 
 // Generates a loop nest that traverse all the unresolved levels in between.
 // TODO: it can only handle all compressed tensors.
if (isCompressedDLT(lvlTypes[tid][firstLvl])) {
                     false,
                     [&](OpBuilder &builder, Location, Value iv,
                         MutableArrayRef<Value> reduc) {
                       ip = builder.saveInsertionPoint();
                       pos = iv;
                       innerArgs.assign(reduc.begin(), reduc.end());
                     })
                     .second;
-            builder.create<scf::YieldOp>(loc, itArgs);
+            YIELD(itArgs);
           });
     } else if (isDenseDLT(lvlTypes[tid][firstLvl])) {
       assert(firstLvl == 0); // This must be the first level.
       Value lb = frontSlice.offset;
       Value sliceSz =
           sliceSizes[tid][*frontSlice.slicedOnLvl][frontSlice.depth - 1];
       Value ub = ADDI(lb, sliceSz);
       outerMost = builder.create<scf::ForOp>(
if (!unResLvls.empty()) {
     // Fills in dense slices levels in between.
     SmallVector<Value> lbs, ubs, steps, lvlSzs;
     for (const SliceInfo *slice : llvm::reverse(unResLvls)) {
       Level sliceLvl = *slice->slicedOnLvl;
       assert(isDenseDLT(lvlTypes[tid][sliceLvl]));
       Value offset = slice->offset;
       Value sliceSz = sliceSizes[tid][sliceLvl][slice->depth - 1];
       lbs.push_back(offset);
-      ubs.push_back(builder.create<arith::AddIOp>(loc, offset, sliceSz));
+      ubs.push_back(ADDI(offset, sliceSz));
       steps.push_back(c1);
       lvlSzs.push_back(lvlSizes[tid][sliceLvl]);
     }
-    auto denseNest = scf::buildLoopNest(
-        builder, loc, lbs, ubs, steps, innerArgs,
-        [&innerArgs, &lvlSzs, &pos,
-         bodyBuilder](OpBuilder &builder, Location loc, ValueRange ivs,
+    auto denseNest =
+        scf::buildLoopNest(builder, loc, lbs, ubs, steps, innerArgs,
+                           [&innerArgs, &lvlSzs, &pos, bodyBuilder](
+                               OpBuilder &builder, Location loc, ValueRange ivs,
                                ValueRange iterArgs) -> scf::ValueVector {
                              for (auto em : llvm::enumerate(ivs)) {
-            // Linearizes postion: pos = (pos * lvlsize) + iv;
-            pos = builder.create<arith::MulIOp>(loc, pos, lvlSzs[em.index()]);
-            pos = builder.create<arith::AddIOp>(loc, pos, em.value());
+                               // Linearizes postion: pos = (pos * lvlsize) +
+                               // iv;
+                               pos = MULI(pos, lvlSzs[em.index()]);
+                               pos = ADDI(pos, em.value());
                              }
                              innerArgs.assign(iterArgs.begin(), iterArgs.end());
                              // Generates user request loop body.
-          // TODO: we do not have to check inbound for dense levels
                              bodyBuilder(builder, loc, pos, innerArgs);
                              return innerArgs;
                            });
-    builder.create<scf::YieldOp>(loc, denseNest.results);
+    YIELD(denseNest.results);
   } else {
     // Generates user request loop body.
     bodyBuilder(builder, loc, pos, innerArgs);
-    builder.create<scf::YieldOp>(loc, innerArgs);
+    YIELD(innerArgs);
   }
   // Insert after current while operation.
   builder.setInsertionPointAfter(outerMost);
   return outerMost.getResults();
 }
 
 void LoopEmitter::genResolvedSliceBegin(OpBuilder &builder, Location loc,
                                         TensorId tid, Level lvl) {
ValueRange result = genUnResolvedSliceTreeTraverse(
       builder, loc, tid, unResSlices, reduc,
       [this, c1, c2, tid, lvl, sPtrBuf](OpBuilder &builder, Location loc,
                                         Value iv,
                                         MutableArrayRef<Value> reduc) {
         Value &nonEmpty = reduc[0];
         Value &minCrd = reduc[1];
         Value &curMemSz = reduc[2];
 
-        Value pHi = builder.create<arith::AddIOp>(loc, iv, c1);
+        Value pHi = ADDI(iv, c1);
         Value sPLo = genIndexLoad(builder, loc, positionsBuffers[tid][lvl], iv);
         Value sPHi =
             genIndexLoad(builder, loc, positionsBuffers[tid][lvl], pHi);
 
         // isNonEmpty = isNonEmpty || lvlNonEmpty, i.e., as long as there is one
         // non-empty lvl, the slice is non-empty.
         Value lvlNonEmpty = CMPI(ult, sPLo, sPHi);
         nonEmpty = builder.create<arith::OrIOp>(loc, lvlNonEmpty, nonEmpty);
 
         // Update the minimum coordinate.
         auto ifNonEmpty = builder.create<scf::IfOp>(loc, builder.getIndexType(),
                                                     lvlNonEmpty, true);
         {
           // Generate Code as follows.
           //
           // if (nonEmpty) {
           //   minCrd = min(minCrd, crd[pos[pLo]]);
           // }
           OpBuilder::InsertionGuard guard(builder);
           builder.setInsertionPointToStart(ifNonEmpty.thenBlock());
           Value curC =
               genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], sPLo);
           Value isSmaller = CMPI(ult, curC, minCrd);
-          Value newMin =
-              builder.create<arith::SelectOp>(loc, isSmaller, curC, minCrd);
-          builder.create<scf::YieldOp>(loc, newMin);
+          Value newMin = SELECT(isSmaller, curC, minCrd);
+          YIELD(newMin);
           builder.setInsertionPointToStart(ifNonEmpty.elseBlock());
-          builder.create<scf::YieldOp>(loc, minCrd);
+          YIELD(minCrd);
         }
         minCrd = ifNonEmpty.getResult(0);
         builder.create<memref::StoreOp>(loc, sPLo, sPtrBuf, curMemSz);
-        Value nxtMemSize = builder.create<arith::AddIOp>(loc, curMemSz, c1);
+        Value nxtMemSize = ADDI(curMemSz, c1);
         builder.create<memref::StoreOp>(loc, sPHi, sPtrBuf, nxtMemSize);
-        // updates the size of the memory curMemSize += 2
-        curMemSz = builder.create<arith::AddIOp>(loc, curMemSz, c2);
+        // curMemSize += 2
+        curMemSz = ADDI(curMemSz, c2);
       });
 
   Value isNonEmpty = result[0];
   Value minCrd = result[1];
   // Two metadata [memSize, idx].
   // TODO: Can use an SSA value for these two metadata
   builder.create<memref::StoreOp>(loc, result[2], sPtrBuf, c0);
   builder.create<memref::StoreOp>(loc, c0, sPtrBuf, c1);
if (depFullyReduced(tid, lvl)) {
     // fully reduced.
     if (isDenseDLT(lvlTypes[tid][lvl]))
       return true;
     // If constraints on the tensor is fully resolved. We do not need to
     // generates slice begin any more, instead we fall back to TACO-based
     // algorithm to (co)iterates over the slice.
     Value pLoPtr =
         genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), c1);
-    pLoPtr = builder.create<arith::AddIOp>(loc, pLoPtr, c2);
-    Value pHiPtr = builder.create<arith::AddIOp>(loc, pLoPtr, c1);
+    pLoPtr = ADDI(pLoPtr, c2);
+    Value pHiPtr = ADDI(pLoPtr, c1);
     posits[tid][lvl] =
         genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), pLoPtr);
     highs[tid][lvl] =
         genIndexLoad(builder, loc, slicePosBuffer[tid][lvl].back(), pHiPtr);
     return true;
   }
 
   // Only when the level is sorted, the next-non-empty slice can be computed
if (slicePosBuffer[tid][lvl][0] == nullptr && !isDenseDLT(lvlType)) {
     // inside the loops. But that would also requires us to allocate/free
     // memorys in loops.
     // TODO: Maybe using allocaScopeOp inside the loop to resolve the issue?
     for (Level curLevel = lvl;
          curLevel >= 1 && !lvlFullyResolved(tid, curLevel - 1); curLevel--) {
       auto depth = remDepOnLevel(tid, curLevel - 1);
       assert(sliceSizes[tid][lvl].size() >= depth);
       Value sz = *(sliceSizes[tid][lvl].rbegin() + depth - 1);
-      bufSize = builder.create<arith::MulIOp>(loc, bufSize, sz);
+      bufSize = MULI(bufSize, sz);
     }
     // For a pair of [pLo, pHi]. Note that we can not compress pHi because slice
     // creates segments in the index buffer so that the pHi for the current
     // level is no longer the pLo for the next level.
-    bufSize = builder.create<arith::MulIOp>(loc, bufSize, c2);
+    bufSize = MULI(bufSize, c2);
     // Additional two metadata {memSize, idx} at head.
-    bufSize = builder.create<arith::AddIOp>(loc, bufSize, c2);
+    bufSize = ADDI(bufSize, c2);
     llvm::for_each(
         slicePosBuffer[tid][lvl], [bufSize, loc, &builder](Value &cache) {
           cache = genAlloca(builder, loc, bufSize, builder.getIndexType());
         });
   }
 
   if (sliceInfo.isInitialTensor() ||
       (lvl >= 1 && lvlFullyResolved(tid, lvl - 1))) {
auto ifOp = builder.create<scf::IfOp>(loc, ValueRange(reduc).getTypes(),
                                         fastPathP, true);
   {
     OpBuilder::InsertionGuard guard(builder);
     // Take the fast path
     // if (minCrd > offset) {
     //   return offset += 1
     // }
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
-    reduc[2] = builder.create<arith::AddIOp>(loc, absOffset, c1);
+    reduc[2] = ADDI(absOffset, c1);
     // Yield offset + 1.
-    builder.create<scf::YieldOp>(loc, reduc);
+    YIELD(reduc);
 
     // else /*minCrd == offset*/ {
     //    for (i = 0; i < slicePos.size(); i+=2) {
     //       if (crd[pos[slicePos[i]]] == minCrd) {
     //          slicePos[i]++;
     //       }
     //       minCrd=min(minCrd, crd[pos[slicePos[i]]]);
     //    }
auto forOp = scf::buildLoopNest(
          &info](OpBuilder &builder, Location loc, ValueRange ivs,
                 ValueRange iterArgs) -> scf::ValueVector {
           Value curMinCrd = iterArgs[0];
           Value isNonEmpty = iterArgs[1];
 
           Type idxTp = builder.getIndexType();
           Value pLo = genIndexLoad(builder, loc, sPtrBuf, ivs.front());
           Value pHi =
-              genIndexLoad(builder, loc, sPtrBuf,
-                           builder.create<arith::AddIOp>(loc, ivs.front(), c1));
+              genIndexLoad(builder, loc, sPtrBuf, ADDI(ivs.front(), c1));
           //
           // if (pLo < pHi) // Only loads when inbound.
           //   coord = load[pLo]
           //   if coord == minCrd
           //     pLo += 1
           //
           // if (pLo < pHi)
           //   curMinCrd = min(curMinCrd, load[pLo])
           //
           Value pred = CMPI(ult, pLo, pHi);
           auto advPLo = builder.create<scf::IfOp>(loc, idxTp, pred, true);
           /* if pLo < pHi */ {
             builder.setInsertionPointToStart(&advPLo.getThenRegion().front());
             // coord = load[pLo]
             Value coord =
                 genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], pLo);
             Value pred = CMPI(eq, coord, info.minCrd);
             auto ifEqual = builder.create<scf::IfOp>(loc, idxTp, pred, true);
             /* if coord == minCrd */ {
               builder.setInsertionPointToStart(
                   &ifEqual.getThenRegion().front());
-              // pLo += 1.
-              Value newPlo = builder.create<arith::AddIOp>(loc, pLo, c1);
+              Value newPlo = ADDI(pLo, c1);
+              // Updates the cache.
               builder.create<memref::StoreOp>(loc, newPlo, sPtrBuf,
                                               ivs.front());
-              builder.create<scf::YieldOp>(loc, newPlo);
+              YIELD(newPlo);
             }
             /* else coord != minCrd */ {
               builder.setInsertionPointToStart(
                   &ifEqual.getElseRegion().front());
-              builder.create<scf::YieldOp>(loc, pLo);
+              YIELD(pLo);
             }
             builder.setInsertionPointAfter(ifEqual);
-            builder.create<scf::YieldOp>(loc, ifEqual.getResults());
+            YIELD(ifEqual.getResults());
           }
           /* else pLo >= pHi */ {
             builder.setInsertionPointToStart(&advPLo.getElseRegion().front());
-            builder.create<scf::YieldOp>(loc, pLo);
+            YIELD(pLo);
           }
 
           builder.setInsertionPointAfter(advPLo);
           pLo = advPLo.getResult(0);
           Value lvlNonEmpty = CMPI(ult, pLo, pHi);
           // Update minCrds
           auto newMin =
               builder.create<scf::IfOp>(loc, idxTp, lvlNonEmpty, true);
           builder.setInsertionPointToStart(&newMin.getThenRegion().front());
-          builder.create<scf::YieldOp>(
-              loc,
-              genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], pLo));
+          YIELD(genIndexLoad(builder, loc, coordinatesBuffers[tid][lvl], pLo));
 
           builder.setInsertionPointToStart(&newMin.getElseRegion().front());
-          builder.create<scf::YieldOp>(loc, curMinCrd);
+          YIELD(curMinCrd);
           builder.setInsertionPointAfter(newMin);
 
           // isNonEmpty = isNonEmpty || lvlNonEmpty
           isNonEmpty =
               builder.create<arith::OrIOp>(loc, lvlNonEmpty, isNonEmpty);
           curMinCrd = builder.create<arith::SelectOp>(
               loc, CMPI(ult, newMin.getResult(0), curMinCrd),
               newMin.getResult(0), curMinCrd);
           return {curMinCrd, isNonEmpty};
         });
 
     builder.setInsertionPointAfter(forOp.loops.front());
     // minOffset = minCrd + 1 >= size ? minCrd + 1 - size : c0
-    Value tmp = builder.create<arith::AddIOp>(loc, forOp.results.front(), c1);
-    Value minOffset = builder.create<arith::SubIOp>(
-        loc, tmp, sliceSizes[tid][lvl][info.depth - 1]);
+    Value tmp = ADDI(forOp.results.front(), c1);
+    Value minOffset = SUBI(tmp, sliceSizes[tid][lvl][info.depth - 1]);
     Value p = CMPI(uge, tmp, sliceSizes[tid][lvl][info.depth - 1]);
-    minOffset = builder.create<arith::SelectOp>(loc, p, minOffset, c0);
+    minOffset = SELECT(p, minOffset, c0);
     SmallVector<Value, 3> yields;
     yields.assign(forOp.results.begin(), forOp.results.end());
     yields.push_back(minOffset);
-    builder.create<scf::YieldOp>(loc, yields);
+    YIELD(yields);
   }
 
   Value nextMinCrd = ifOp.getResults()[0];
   Value nextNonEmpty = ifOp.getResults()[1];
 
   // The next offset should at least be offset + 1;
   Value minOffset = ifOp.getResults()[2];
-  Value nxOffset = builder.create<arith::AddIOp>(loc, info.offset, c1);
+  Value nxOffset = ADDI(info.offset, c1);
   Value maxPred = CMPI(ugt, minOffset, nxOffset);
-  Value nextAbsOffset =
-      builder.create<arith::SelectOp>(loc, maxPred, minOffset, nxOffset);
+  Value nextAbsOffset = SELECT(maxPred, minOffset, nxOffset);
 
-  Value sliceUB = builder.create<arith::AddIOp>(
-      loc, nextAbsOffset, sliceSizes[tid][lvl][info.depth - 1]);
+  Value sliceUB = ADDI(nextAbsOffset, sliceSizes[tid][lvl][info.depth - 1]);
 
   // FIXME: this only works if there is only one parent.
   assert(info.depth - 1 == 0);
   // nextNonEmpty = nextNonEmpty && slice upper bound <= parent upperbound.
-  nextNonEmpty = builder.create<arith::AndIOp>(
-      loc, nextNonEmpty, CMPI(ule, sliceUB, lvlSizes[tid][lvl]));
+  nextNonEmpty = ANDI(nextNonEmpty, CMPI(ule, sliceUB, lvlSizes[tid][lvl]));
 
   // FIXME: compute relative offset.
   assert(info.depth - 1 == 0);
   Value nextRelOffset = nextAbsOffset;
-  nextRelOffset =
-      builder.create<arith::SelectOp>(loc, nextNonEmpty, nextRelOffset, c0);
+  nextRelOffset = SELECT(nextNonEmpty, nextRelOffset, c0);
 
   operands.push_back(nextNonEmpty);
   operands.push_back(nextMinCrd);
   operands.push_back(nextAbsOffset); // we push the absolute offset.
 
   // Update the slice stack.
   info.isNonEmpty = whileOp.getResult(retIdx++);
   info.minCrd = whileOp.getResult(retIdx++);
auto whileOp = builder.create<scf::WhileOp>(
       });
 
   // Set the insertion point to while loop body.
   builder.setInsertionPointToEnd(&whileOp.getAfter().front());
   return whileOp;
 }
 
 #undef CMPI
-#undef ADDI
 #undef C_IDX
+#undef YIELD
+#undef ADDI
+#undef ANDI
+#undef SUBI
+#undef MULI
+#undef SELECT