ArithIndexingBuilder

Author: Hardcode84

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

@@ -101,7 +101,10 @@ Type getType(OpFoldResult ofr);
 /// Helper struct to build simple arithmetic quantities with minimal type
 /// inference support.
 struct ArithBuilder {
-  ArithBuilder(OpBuilder &b, Location loc) : b(b), loc(loc) {}
+  ArithBuilder(
+      OpBuilder &b, Location loc,
+      arith::IntegerOverflowFlags ovf = arith::IntegerOverflowFlags::none)
+      : b(b), loc(loc), ovf(ovf) {}
 
   Value _and(Value lhs, Value rhs);
   Value add(Value lhs, Value rhs);
@@ -114,6 +117,15 @@ struct ArithBuilder {
 private:
   OpBuilder &b;
   Location loc;
+  arith::IntegerOverflowFlags ovf;
+};
+
+/// ArithBuilder specialized specifically for tensor/memref indexing
+/// calculations. Those calculations generally should never signed overflow, so
+/// we can set oveflow flags accordingly.
+struct ArithIndexingBuilder : public ArithBuilder {
+  ArithIndexingBuilder(OpBuilder &b, Location loc)
+      : ArithBuilder(b, loc, arith::IntegerOverflowFlags::nsw) {}
 };
 
 namespace arith {

mlir/lib/Dialect/Arith/Utils/Utils.cpp

@@ -315,17 +315,17 @@ Value ArithBuilder::_and(Value lhs, Value rhs) {
 Value ArithBuilder::add(Value lhs, Value rhs) {
   if (isa<FloatType>(lhs.getType()))
     return b.create<arith::AddFOp>(loc, lhs, rhs);
-  return b.create<arith::AddIOp>(loc, lhs, rhs);
+  return b.create<arith::AddIOp>(loc, lhs, rhs, ovf);
 }
 Value ArithBuilder::sub(Value lhs, Value rhs) {
   if (isa<FloatType>(lhs.getType()))
     return b.create<arith::SubFOp>(loc, lhs, rhs);
-  return b.create<arith::SubIOp>(loc, lhs, rhs);
+  return b.create<arith::SubIOp>(loc, lhs, rhs, ovf);
 }
 Value ArithBuilder::mul(Value lhs, Value rhs) {
   if (isa<FloatType>(lhs.getType()))
     return b.create<arith::MulFOp>(loc, lhs, rhs);
-  return b.create<arith::MulIOp>(loc, lhs, rhs);
+  return b.create<arith::MulIOp>(loc, lhs, rhs, ovf);
 }
 Value ArithBuilder::sgt(Value lhs, Value rhs) {
   if (isa<FloatType>(lhs.getType()))

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

@@ -1102,15 +1102,14 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(loadOp);
     Location loc = loadOp.getLoc();
+    ArithIndexingBuilder idxBuilderf(rewriter, loc);
     for (auto i : llvm::seq<int64_t>(rankOffset, indices.size() - finalRank)) {
       OpFoldResult pos = extractPos[i - rankOffset];
       if (isConstantIntValue(pos, 0))
         continue;
 
       Value offset = getValueOrCreateConstantIndexOp(rewriter, loc, pos);
-
-      auto ovf = arith::IntegerOverflowFlags::nsw;
-      indices[i] = rewriter.create<arith::AddIOp>(loc, indices[i], offset, ovf);
+      indices[i] = idxBuilderf.add(indices[i], offset);
     }
 
     Value base = loadOp.getBase();