Merge pull request #378 from Arech8:pr0_improve_inout

PiperOrigin-RevId: 881037610

Author: Google-ML-Automation

README.md

@@ -22,13 +22,13 @@ import triton.language as tl
 
 @triton.jit
 def add_kernel(
-    x_ptr,
-    y_ptr,
-    length,
-    output_ptr,
-    block_size: tl.constexpr,
+    x_ptr,        # First 3 arguments
+    y_ptr,        # are input
+    length,       # arguments.
+    output_ptr,   # Implicit output argument goes after inputs.
+    block_size: tl.constexpr, # Constexpr params goes the last.
 ):
-  """Adds two vectors."""
+  """Adds two vectors output = x + y."""
   pid = tl.program_id(axis=0)
   block_start = pid * block_size
   offsets = block_start + tl.arange(0, block_size)
@@ -47,43 +47,87 @@ import jax.numpy as jnp
 import jax_triton as jt
 
 def add(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
-  out_shape = jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype)
   block_size = 8
   return jt.triton_call(
-      x,
-      y,
-      x.size,
+      x,                  # Kernel's input arguments are the first
+      y,                  # in jt.triton_call(). The output argument
+      x.size,             # is passed implicitly.
       kernel=add_kernel,
-      out_shape=out_shape,
+      out_shape=x,
       grid=(x.size // block_size,),
-      block_size=block_size)
+      block_size=block_size   # Constexpr params are passed as kwargs
+    )
 
 x_val = jnp.arange(8)
 y_val = jnp.arange(8, 16)
 print(add(x_val, y_val))
 print(jax.jit(add)(x_val, y_val))
 ```
 
+One could also use input-output parameters for kernels:
+
+```python
+
+@triton.jit
+def add_inplace_y_kernel(
+    x_ptr,          # input vector
+    y_inout_ptr,    # explicit in-out vector (could be anywhere)
+    length,
+    block_size: tl.constexpr,
+):
+  """Adds two vectors output = x + y."""
+  pid = tl.program_id(axis=0)
+  block_start = pid * block_size
+  offsets = block_start + tl.arange(0, block_size)
+  mask = offsets < length
+  x = tl.load(x_ptr + offsets, mask=mask)
+  y = tl.load(y_inout_ptr + offsets, mask=mask)
+  output = x + y
+  tl.store(y_inout_ptr + offsets, output, mask=mask)
+
+
+from functools import partial
+
+# jitting or jitting with donation isn't mandatory, but makes invocation more efficient.
+# Otherwise XLA would have to make a copy of each non-donated in-out argument before
+# calling a kernel, since JAX arrays by default are immutable.
+@partial(jax.jit, donate_argnames="y")
+def add_inplace_y(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:
+  block_size = 8
+  return jt.triton_call(
+      x,
+      y,            # explicit in-out argument
+      x.size,
+      kernel=add_inplace_y_kernel,
+      input_output_aliases={1: 0},  # input arg idx 1 (y) is the first output arg
+      out_shape=x,
+      grid=(x.size // block_size,),
+      block_size=block_size)
+
+x_val = jnp.arange(8)
+y_val = jnp.arange(8, 16)
+print(add_inplace_y(x_val, y_val))
+```
+
 See [the examples
 directory](https://github.com/jax-ml/jax-triton/tree/main/examples), especially
 [fused_attention.py](https://github.com/jax-ml/jax-triton/blob/main/examples/fused_attention.py)
 and [the fused attention
 ipynb](https://github.com/jax-ml/jax-triton/blob/main/examples/JAX_%2B_Triton_Flash_Attention.ipynb).
 
+Some other use-cases are also covered in [tests](https://github.com/jax-ml/jax-triton/tree/main/tests).
+
 ## Installation
 
 ```bash
 $ pip install jax-triton
 ```
 
-Make sure you have a CUDA-compatible `jax` installed. For example you could run:
+Make sure you have a CUDA- or ROCm- compatible `jax` installed. For example you could run:
 ```bash
 $ pip install "jax[cuda12]"
 ```
 
-`jax-triton` currently requires building the latest version of `triton`
-[from source](https://triton-lang.org/main/getting-started/installation.html#from-source).
-
 ## Development
 
 To develop `jax-triton`, you can clone the repo with:

jax_triton/triton_lib.py

@@ -34,7 +34,6 @@
 from jax._src import state
 from jax._src import util
 from jax._src.lib import gpu_triton as triton_kernel_call_lib
-import jax.dlpack
 import jax.extend as jex
 from jax.interpreters import ad
 from jax.interpreters import batching
@@ -435,6 +434,19 @@ def get_or_create_triton_kernel(
   kernel = _COMPILED_KERNEL_CACHE.get(cache_key)
 
   if kernel is None:
+    # First, check that the kernel signature and the reconstructed signature have the
+    # same number of parameters. A mismatch can occur due to differences in
+    # `triton_call(input_output_aliases=)` handling between jax-triton versions.
+    if len(fn.signature.parameters) != len(signature):
+      raise TypeError(
+        f"Number of parameters in the kernel '{fn}' signature "
+        f"({len(fn.signature.parameters)}: {fn.signature}) "
+        f"does not match reconstructed signature ({len(signature)}: {signature}). "
+        "If the kernel was working on an older version of jax-triton and its "
+        "triton_call() launcher uses `input_output_aliases` argument, note that "
+        "implicit output arguments are no longer required for aliased args."
+      )
+
     opts = {
         "num_warps": num_warps,
         "num_stages": num_stages,
@@ -543,8 +555,17 @@ def triton_kernel_call_lowering(
   for idx, dtype, v in scalar_args:
     args.insert(idx, v)
     arg_dtypes.insert(idx, dtype)
-  args.extend(ctx.avals_out)
-  arg_dtypes.extend(map(get_triton_type, ctx.avals_out))
+  # Extract only the output avals not referenced in the input_output_aliases mapping.
+  assert isinstance(input_output_aliases, tuple)
+  input_output_aliases = dict(input_output_aliases)
+  strictly_out_avals = [
+    aval
+    for i, aval in enumerate(ctx.avals_out)
+    if i not in input_output_aliases.values()
+  ]
+  args.extend(strictly_out_avals)
+  arg_dtypes.extend(map(get_triton_type, strictly_out_avals))
+
   named_args = dict(unsafe_zip(fn.arg_names, args))
 
   if isinstance(fn, autotuner.Autotuner):
@@ -606,6 +627,10 @@ def prune_configs(configs, named_args, **kwargs):
         "`kernel` must be a Triton `JITFunction`, `Heuristics` or `Autotuner`."
     )
 
+  output2input = {v: k for k, v in input_output_aliases.items()}
+  if len(output2input) != len(input_output_aliases):
+    raise ValueError("input_output_aliases must be a bijection")
+
   outputs_offset = len(ctx.avals_in) + len(scalar_args)
   config_params = []
   for config in configs:
@@ -616,9 +641,13 @@ def prune_configs(configs, named_args, **kwargs):
     if callable(zeroed_outputs):
       config_zeroed_outputs = config_zeroed_outputs(config_metaparams)
 
+    # zeroed_params_with_sizes is a dict output_arg_idx -> aval_size_bytes
+    # config_zeroed_outputs is output ordinal numbers
     zeroed_params_with_sizes = {
-        i + outputs_offset: aval_size_bytes(ctx.avals_out[i])
-        for i in sorted(config_zeroed_outputs)
+      output2input[i] if i in output2input else i + outputs_offset: aval_size_bytes(
+        ctx.avals_out[i]
+      )
+      for i in sorted(config_zeroed_outputs)
     }
 
     config_params.append(
@@ -688,7 +717,7 @@ def prune_configs(configs, named_args, **kwargs):
     named_scalar_args = {fn.arg_names[i]: v for i, _, v in scalar_args}
     input_output_aliases_with_sizes = tuple(
         (input_idx, output_idx, aval_size_bytes(ctx.avals_in[input_idx]))
-        for input_idx, output_idx in input_output_aliases
+        for input_idx, output_idx in input_output_aliases.items()
     )
     kernel_call = triton_kernel_call_lib.TritonAutotunedKernelCall(
         f"{kernel_call_name} ({fn.fn.__name__}) {named_scalar_args}",
@@ -703,7 +732,7 @@ def prune_configs(configs, named_args, **kwargs):
       custom_call_target_name,
       api_version=2,
       backend_config=zlib.compress(call_proto),
-      operand_output_aliases=dict(input_output_aliases),
+      operand_output_aliases=input_output_aliases,
   )
   return rule(ctx, *array_args)