[Mlir-commits] [mlir] [mlir][scf]: Avoid using wrong calculation loop pipelining kernel upperbound (PR #116748)

Aviad Cohen llvmlistbot at llvm.org
Mon Nov 18 23:46:35 PST 2024


================
@@ -444,7 +444,15 @@ scf::ForOp LoopPipelinerInternal::createKernelLoop(
         loc, rewriter.getIntegerAttr(t, maxStage));
     Value maxStageByStep =
         rewriter.create<arith::MulIOp>(loc, step, maxStageValue);
-    newUb = rewriter.create<arith::SubIOp>(loc, ub, maxStageByStep);
+    Value hasAtLeastOneIteration = rewriter.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::slt, maxStageByStep, ub);
+    Value possibleNewUB =
+        rewriter.create<arith::SubIOp>(loc, ub, maxStageByStep);
+    // In case of `index` or `unsigned` type, we need to make sure that the
+    // subtraction does not result in a negative value, instead we use lb
+    // to avoid entering the kernel loop.
+    newUb = rewriter.create<arith::SelectOp>(
+        loc, hasAtLeastOneIteration, possibleNewUB, forOp.getLowerBound());
----------------
AviadCo wrote:

@ThomasRaoux , actualy what we want to make sure is that maxStageByStep < ub, because if those value are from `unsigned` or `index` may lead to wrong answer.
For exmaple, if:
ub = 100
maxStageByStep = 150
Then:
100 - 150 = (-50) 
But for index/unsigned this value will be intreprated to **big possitive number** as negative doesn't exist in unsigned.
This leads to error that we will actually enter the kernel loop although `ub < maxStageByStep`.
So in that case, to avoid using `ub-maxStageByStep < lb` which in the case mentioned above may be falsely true, I added a check and use `lb` just to avoid entering the kernel.
The rest of the code using `ub-maxStageByStep` is protected by `pred` condition and usually wrapped with `if` on the `pred` so I think we are o.k. there, the only place where it didn't work for me was on the `for` kernel mentioned above.


https://github.com/llvm/llvm-project/pull/116748


More information about the Mlir-commits mailing list