Add IndexCastUIOp to AxisInfo.

`index_cast` is supported, so there is not reason why `index_castui` shouldn't be supported. `index_castui` happens to appear in our Triton fusion emitter. Without this change Triton fails to correctly propagate axis info and do layout optimization. PiperOrigin-RevId: 696500599
openxla · Nov 14, 2024 · 2465051 · 2465051
1 parent f908c44
commit 2465051
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/third_party/triton/temporary/index_cast_ui_axis_info.patch b/third_party/triton/temporary/index_cast_ui_axis_info.patch
@@ -0,0 +1,10 @@
+--- a/lib/Analysis/AxisInfo.cpp
++++ b/lib/Analysis/AxisInfo.cpp
+@@ -1011,6 +1011,7 @@ AxisInfoAnalysis::AxisInfoAnalysis(DataF
+                   CastOpAxisInfoVisitor<arith::ExtUIOp>,
+                   CastOpAxisInfoVisitor<arith::TruncIOp>,
+                   CastOpAxisInfoVisitor<arith::IndexCastOp>,
++                  CastOpAxisInfoVisitor<arith::IndexCastUIOp>,
+                   CastOpAxisInfoVisitor<triton::gpu::ConvertLayoutOp>,
+                   CastOpAxisInfoVisitor<mlir::UnrealizedConversionCastOp>,
+                   CastOpAxisInfoVisitor<triton::BitcastOp>>();
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
@@ -16,5 +16,6 @@ those to this list.
 temporary_patch_list = [
     "//third_party/triton:temporary/replace_unreachable_by_abort.patch",
     "//third_party/triton:temporary/block_k_16_fix.patch",
+    "//third_party/triton:temporary/index_cast_ui_axis_info.patch",
     # Add new patches just above this line
 ]