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

[aartbik]

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

36–38

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 ADDI(lhs, rhs) (builder.create<arith::AddIOp>(loc, (lhs), (rhs)))
+#define MULI(lhs, rhs) (builder.create<arith::MulIOp>(loc, (lhs), (rhs)))
 
-#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. The generated code returns
 /// constant 0 if `isNonEmpty` is false.
 ///
 /// 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);
   // offset
-  Value mp1 = builder.create<arith::AddIOp>(loc, minCrd, C_IDX(1));
+  Value mp1 = ADDI(minCrd, C_IDX(1));
   Value mms = builder.create<arith::SubIOp>(loc, mp1, size);
   // This is the absolute offset related to the underly tensor.
   return builder.create<arith::SelectOp>(loc, 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

 
 //===----------------------------------------------------------------------===//
 // 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>(
           builder.setInsertionPointToStart(ifInBound.elseBlock());
           builder.create<scf::YieldOp>(loc, 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));
+        Value nextPos = ADDI(ivs[0], C_IDX(1));
         builder.create<scf::YieldOp>(loc, 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()) {
         // Needs at least two operands to form a non-trivial affine expression.
         ArrayRef<std::pair<TensorId, Level>> depLvls = dependentLvlMap[t][lvl];
         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 {
       Value c1 = C_IDX(1), c2 = C_IDX(2);
       if (!isDenseDLT(lvlTypes[tid][lvl])) {
         // 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.
         // TODO: we could probably use an SSA value for it.
         Value sPtrBuf = slicePosBuffer[tid][lvl].back();
         Value curP = genIndexLoad(builder, loc, sPtrBuf, c1);
-        Value nexP = builder.create<arith::AddIOp>(loc, curP, c2);
+        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");
   }
if (isDenseSliceCond) {
       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 = 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)) {
     // The isSparseSliceCond is handled properly already.
     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);
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));
+      Value nextPos = ADDI(yields[idx], C_IDX(1));
       yields[idx] =
           builder.create<arith::SelectOp>(loc, 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);
     }
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));
       // 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!
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);
   builder.setInsertionPointAfter(whileOp);
 }
size_t LoopEmitter::remDepOnLevel(TensorId tid, Level lvl) const {
   return totalDependencies;
 }
 
 std::pair<Operation *, ValueRange> LoopEmitter::genSliceLvlTraverseLoop(
     OpBuilder &builder, Location loc, Value loopLo, Value loopHi, Value offset,
     TensorId tid, Level lvl, size_t depth, ValueRange userReduc, bool genYield,
     LoopBodyBuilder bodyBuilder) {
   Value c1 = C_IDX(1);
-  Value sliceHi =
-      builder.create<arith::AddIOp>(loc, offset, sliceSizes[tid][lvl].back());
+  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>(
scf::WhileOp whileOp = builder.create<scf::WhileOp>(
           }
         }
         // 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));
+        yields.push_back(ADDI(iv, c1));
         yields.push_back(cont);
         yields.append(ifOp.getResults().begin(), ifOp.getResults().end());
         builder.create<scf::YieldOp>(loc, yields);
       });
 
   builder.setInsertionPointAfter(whileOp);
   return std::make_pair(whileOp, whileOp.getResults().drop_front(2));
 }
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, &ip,
-         &pos](OpBuilder &builder, Location loc, ValueRange ivs,
+    auto denseNest =
+        scf::buildLoopNest(builder, loc, lbs, ubs, steps, innerArgs,
+                           [&innerArgs, &lvlSzs, &ip, &pos](
+                               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());
                              }
                              ip = builder.saveInsertionPoint();
                              innerArgs.assign(iterArgs.begin(), iterArgs.end());
                              return innerArgs;
                            });
   }
   // Generates user request loop body.
   builder.restoreInsertionPoint(ip);
   bodyBuilder(builder, loc, pos, innerArgs);
   builder.create<scf::YieldOp>(loc, innerArgs);
 
   // Insert after current while operation.
   builder.setInsertionPointAfter(outterMost);
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
         Value lvlNonEmpty = CMPI(ult, sPLo, sPHi);
         nonEmpty = builder.create<arith::OrIOp>(loc, lvlNonEmpty, nonEmpty);
 
ValueRange result = genUnResolvedSliceTreeTraverse(
               builder.create<arith::SelectOp>(loc, isSmaller, curC, minCrd);
           builder.create<scf::YieldOp>(loc, newMin);
           builder.setInsertionPointToStart(ifNonEmpty.elseBlock());
           builder.create<scf::YieldOp>(loc, minCrd);
         }
         minCrd = ifNonEmpty.getResult(0);
         // filles in
         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);
         // curMemSize += 2
-        curMemSz = builder.create<arith::AddIOp>(loc, curMemSz, c2);
+        curMemSz = ADDI(curMemSz, c2);
       });
 
   unsigned depth = levelReducedDep[tid][lvl];
   Value size = sliceSizes[tid][lvl][depth];
   Value isNonEmpty = result[0];
   Value minCrd = result[1];
   // Two metadata [memSize, idx].
   // TODO: Can use an SSA value for these two metadata
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))) {
void LoopEmitter::genSliceNextInduction(OpBuilder &builder, Location loc,
   Value sPtrBuf = slicePosBuffer[tid][lvl][info.depth - 1];
   Value fastPathP = CMPI(ugt, info.minCrd, absOffset);
   auto ifOp = builder.create<scf::IfOp>(loc, ValueRange(reduc).getTypes(),
                                         fastPathP, true);
   {
     OpBuilder::InsertionGuard guard(builder);
     // Take the fast path if minCrd > offset
     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);
 
     // Else, take the slow path.
     builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
     reduc[2] = absOffset; // restore value.
     Value pSt = c2;       // pointer starting index
     Value mSz = genIndexLoad(builder, loc, sPtrBuf, c0); // memSize
     reduc[0] = lvlSizes[tid][lvl];                       // next min coord
     reduc[1] = constantI1(builder, loc, false);          // isNonEmpty
     auto loopArgs = static_cast<ValueRange>(reduc).drop_back();
     auto forOp = scf::buildLoopNest(
         builder, loc, pSt, mSz, c2, loopArgs,
         [this, tid, lvl, c1, sPtrBuf,
          &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
           //   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());
-              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);
             }
             /* else coord != minCrd */ {
               builder.setInsertionPointToStart(
                   &ifEqual.getElseRegion().front());
auto forOp = scf::buildLoopNest(
           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 tmp = ADDI(forOp.results.front(), c1);
     Value minOffset = builder.create<arith::SubIOp>(
         loc, 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);
     SmallVector<Value, 3> yields;
     yields.assign(forOp.results.begin(), forOp.results.end());
     yields.push_back(minOffset);
     builder.create<scf::YieldOp>(loc, 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 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 the parsent is the tensor, we should use the
   // parents slice size + parent offset.
   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]));
 
Operation *LoopEmitter::emitSliceDrivenLoopOverTensorAtLvl(
   // levelReducedDep[tid][lvl]++;
   // Set the insertion point to while loop body.
   builder.setInsertionPointToEnd(&whileOp.getAfter().front());
   return whileOp;
 }
 
 #undef CMPI
 #undef ADDI
+#undef MULI
 #undef C_IDX