[Triton] Fixing getLastInductionValue utility to also accept Index type. This would otherwise crash when warp specialization is enabled.

PiperOrigin-RevId: 820159796

Author: Moerafaat

third_party/triton/temporary/series.bzl

@@ -14,5 +14,6 @@ those to this list.
 """
 
 temporary_patch_list = [
+    "//third_party/triton:temporary/utility-fix.patch",
     # Add new patches just above this line
 ]

third_party/triton/temporary/utility-fix.patch

@@ -0,0 +1,22 @@
+This patch would probably not be accepted upstream because our infrastructure
+uses Index type for indexing, while they use Integer type. Triton frontend
+wouldn't generate a situation that would run into this issue.
+
+diff --git a/lib/Dialect/Triton/IR/Utility.cpp b/lib/Dialect/Triton/IR/Utility.cpp
+--- a/lib/Dialect/Triton/IR/Utility.cpp
++++ b/lib/Dialect/Triton/IR/Utility.cpp
+@@ -97,8 +97,12 @@ Value tt::getLastInductionValue(OpBuilde
+   // (ub - lb -1) // step * step + lb
+   Value diff =
+       b.create<arith::SubIOp>(loc, loop.getUpperBound(), loop.getLowerBound());
+-  diff = b.create<arith::SubIOp>(
+-      loc, diff, b.create<arith::ConstantOp>(loc, b.getI32IntegerAttr(1)));
++  Value one;
++  if (diff.getType().isIndex())
++    one = b.create<arith::ConstantIndexOp>(loc, 1);
++  else
++    one = b.create<arith::ConstantOp>(loc, b.getIntegerAttr(diff.getType(), 1));
++  diff = b.create<arith::SubIOp>(loc, diff, one);
+   Value ceilStep = b.create<arith::MulIOp>(
+       loc, b.create<arith::DivSIOp>(loc, diff, loop.getStep()), loop.getStep());
+   return b.create<arith::AddIOp>(loc, ceilStep, loop.getLowerBound());