Background

This PR adds custom op support for __builtin_index_select, whose semantics align with torch.index_select.

This change improves the expressiveness of index-related kernels and enables more general indexing patterns in Triton custom ops.

Interface Introduction

The operator gathers elements from a source tensor in Global Memory (GM) into an output tensor in Unified Buffer (UB) using indices provided in UB along a specified dimension.

Usage Example

@triton.jit
def builtin_index_select_kernel(src_ptr, index_ptr, out_ptr):
    # Define 2x2 tile indices for output tensor
    r = tl.arange(0, 2)[:, None]  # Row indices: shape [2, 1]
    c = tl.arange(0, 2)[None, :]  # Column indices: shape [1, 2]

    # Load index tensor (shape [2]) from GM to UB
    idx = tl.load(index_ptr + tl.arange(0, 2))
    # Initialize empty 2x2 output tile in UB (default value: 0)
    dst = tl.full((2, 2), 0, dtype=tl.float32)

    # Invoke __builtin_index_select custom op to gather elements
    out_tile = al.custom(
        "__builtin_index_select",
        src_ptr,          # Pointer to source tensor in GM
        idx,              # Index tensor (in UB) for gathering
        dim=0,            # Dimension to gather along
        bound=4,          # Upper bound for valid index values (out-of-bound check)
        end_offset=(2, 2),# End offsets of each dimension for the index tensor
        start_offset=(0, 0), # Start offsets of each dimension for the source tensor
        src_stride=(4, 1),# Stride of each dimension for the source tensor in GM
        out=dst           # Output tensor (in UB) to store gathered elements
    )

    # Store the gathered tile from UB to output tensor in GM
    tl.store(out_ptr + r * 2 + c, out_tile)

Checklist

• I am not making a trivial change, such as fixing a typo in a comment.
• I have written a PR description following these
  rules.
• I have run pre-commit run --from-ref origin/main --to-ref HEAD.
• Select one of the following.
  • I have added tests.
    • /test for lit tests
    • /unittest for C++ tests
    • /python/test for end-to-end tests
  • This PR does not need a test because FILL THIS IN.
• Select one of the following.
  • I have not added any lit tests.
  • The lit tests I have added follow these best practices,
    including the "tests should be minimal" section. (Usually running Python code
    and using the instructions it generates is not minimal.)