Integrate Triton up to [bfb8e413](https://github.com/openai/triton/commits/bfb8e413b075583228c961bdbb65a98dc54d0868)

PiperOrigin-RevId: 614740360

Author: gflegar

jax/_src/pallas/triton/lowering.py

@@ -492,7 +492,7 @@ def _associative_scan_lowering(body, ctx: LoweringRuleContext, args, axes):
   if consts:
     raise NotImplementedError("Associative scan with constants not supported.")
   element_types = [_element_type(arg.type) for arg in flat_args]
-  scan_op = tt_dialect.ScanOp(flat_args, axis)
+  scan_op = tt_dialect.ScanOp(flat_args, axis, reverse=False)
   param_types = element_types * 2
   entry = scan_op.regions[0].blocks.append(*param_types)
   with ir.InsertionPoint.at_block_begin(entry):

jaxlib/triton/dialect.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+import inspect
 
 from jaxlib.mlir._mlir_libs._triton_ext import (
     PointerType,
@@ -55,12 +56,18 @@ def __init__(
       self,
       operands: Sequence[ir.Value],
       axis: int,
+      reverse: ir.Attribute,
       *,
       loc: ir.Location | None = None,
       ip: ir.InsertionPoint | None = None,
   ):
     return_types = [op.type for op in operands]
-    super().__init__(return_types, operands, axis, loc=loc, ip=ip)
+    # OSS might have an old version of Triton, whose ScanOp doesn't have a
+    # reverse parameter.
+    if "reverse" in inspect.signature(super().__init__).parameters:
+      super().__init__(return_types, operands, axis, reverse, loc=loc, ip=ip)
+    else:
+      super().__init__(return_types, operands, axis, loc=loc, ip=ip)  # pylint: disable=no-value-for-parameter
 
 
 # TODO(slebedev): Consider overriding instead.