[Bug][vector][mthreads] Vectorization missed due to non-vector usage immediately after vector load.

[liuxiao0909c]

docker: registry.mthreads.com/mcconline/inference/vllm:v0.20.2-ph1-4.3.5-torch2.7.1-v1.1.0

flagtree: 0.6.0+mthreads.gitc64a4918(after the fix of #854)

description: VectorCombinePass replaces load <4 x float> with four load float

testcase:

import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle


@triton.jit
def _vec_kernel(
    x_ptr,
    y_ptr,
    x_stride0,
    y_stride0,
    VEC: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    row_id = tl.program_id(0)
    x_ptr += row_id * x_stride0
    y_ptr += row_id * y_stride0
    vec = tl.arange(0, VEC)
    lane = tl.arange(0, BLOCK_SIZE)
    vals = tl.load(x_ptr + lane[:, None] * VEC + vec[None, :])
    tl.store(y_ptr, tl.sum(vals))  # no vec
    #tl.store(y_ptr + lane[:, None] * VEC + vec[None, :], vals) # vec


device='musa'
torch.manual_seed(42)

VEC = 4
BLOCK_SIZE = 512
num_rows = 2
num_elem = BLOCK_SIZE * VEC
X = torch.randn(num_rows, num_elem, device=device, dtype=torch.float32)
Y = torch.empty((num_rows, num_elem,), dtype=torch.float32, device=device)
_vec_kernel[(num_rows,)](
    X,
    Y,
    X.stride(0),
    Y.stride(0),
    VEC=VEC,
    BLOCK_SIZE=BLOCK_SIZE,
    num_warps=BLOCK_SIZE // 32,
)

IR after SROAPass runs:

mlir after SLPVectorizerPass.txt

IR after VectorCombinePass runs:

mlir after VectorCombinePass.txt

[liuxiao0909c]

additional testcase:

the .ll file of the previous testcase compiled by llc has a recovered load <4 x float>

import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle


@triton.jit
def tle_top_k_per_row_prefill(
    logits_ptr,
    row_starts,
    row_ends,
    stride0,
    vocab_size,
    dump_buffer_ptr,
    BLOCK_SIZE: tl.constexpr,
):
    NUM_BINS: tl.constexpr = 2048
    VEC: tl.constexpr = 4

    row_id = tl.program_id(0)
    row_start = tl.load(row_starts + row_id)
    row_end = tl.load(row_ends + row_id)
    logits_ptr += row_id * stride0
    dump_buffer_ptr += row_id * NUM_BINS
    # float4 align
    x_off_mod = (row_id * stride0 + row_start) % VEC
    skip_elems = 0 if x_off_mod == 0 else VEC - x_off_mod
    # alloc counter
    s_histogram = tle.gpu.alloc(
        [NUM_BINS],
        dtype=tl.int32,
        layout=None,
        scope=tle.gpu.smem,
        nv_mma_shared_layout=False,
    )
    s_histogram_ptr = tle.gpu.local_ptr(s_histogram, (0,))
    # clear counter
    lane = tl.arange(0, BLOCK_SIZE)
    vec = tl.arange(0, VEC)
    threshold_rounds: tl.constexpr = NUM_BINS // BLOCK_SIZE
    for clear_round in tl.static_range(0, threshold_rounds):
        clear_bins = clear_round * BLOCK_SIZE + lane
        tl.store(s_histogram_ptr + clear_bins, 0)
    tl.debug_barrier()
    # statistic
    aligned_row_ptr = tl.multiple_of(logits_ptr + row_start + skip_elems, VEC * 4)
    row_len = row_end - row_start - skip_elems
    n_vec_full = row_len // (BLOCK_SIZE * VEC)
    for t in tl.range(0, n_vec_full):
        base = t * BLOCK_SIZE * VEC + lane * VEC
        offs = base[:, None] + vec[None, :]
        x_vec = tl.load(aligned_row_ptr + offs)
        h = x_vec.to(tl.float16)
        bits = h.to(tl.uint16, bitcast=True)
        sign_mask = tl.full(bits.shape, 0x8000, tl.uint16)
        sign_set = (bits & sign_mask) != 0
        inv = (~bits) & tl.full(bits.shape, 0x7FFF, tl.uint16)
        mapped = tl.where(sign_set, bits, inv)
        bin_idx = (mapped >> 5).to(tl.uint32)
        tl.atomic_add(
            s_histogram_ptr + bin_idx,
            1,
            sem="relaxed",
            scope="cta",
        )
    tl.debug_barrier()
    # copy counter to dump_buffer
    counts = tl.load(s_histogram_ptr + tl.arange(0, NUM_BINS))
    tl.store(dump_buffer_ptr + tl.arange(0, NUM_BINS), counts)


device='musa'
NUM_THREADS_PER_BLOCK = 512
NUM_BINS = 2048

torch.manual_seed(789)
num_rows = 1
vocab_size = 129280
top_k = 1024

logits = torch.randn(num_rows, vocab_size, device=device, dtype=torch.float32)
row_starts = torch.zeros(num_rows, dtype=torch.int32, device=device)
row_ends = torch.full((num_rows,), vocab_size, dtype=torch.int32, device=device)
stride0 = logits.stride(0)
dump_buffer = torch.empty((num_rows, NUM_BINS), device=device, dtype=torch.int32)

tle_top_k_per_row_prefill[(num_rows,)](
    logits,
    row_starts,
    row_ends,
    stride0,
    vocab_size,
    dump_buffer,
    BLOCK_SIZE=NUM_THREADS_PER_BLOCK,
    num_warps=NUM_THREADS_PER_BLOCK // 32,
)
torch.musa.synchronize()

ir after SROAPass：

ir after VectorCombinePass:

ir before llc compile .ll to .o:

[Kylin1207]

This optimization will be offloaded to the MTCC. The resulting binaries are identical regardless of whether VectorCombine is enabled or disabled. I have compared the assembly outputs of both versions, and the MTCC compiles it down to a single vectorized load instruction for <4xfloat>

[liuxiao0909c]

This optimization will be offloaded to the MTCC. The resulting binaries are identical regardless of whether VectorCombine is enabled or disabled. I have compared the assembly outputs of both versions, and the MTCC compiles it down to a single vectorized load instruction for <4xfloat>

Thank you for your reply! I observed two command invocation sequences during Triton MUSA backend compilation:

['/usr/local/lib/python3.10/dist-packages/triton/backends/mthreads/bin/llc', '/tmp/triton-musa-8y4br0gb/kernel.ll', '-march=mtgpu', '-mcpu=mp_31', '--opaque-pointers', '-O2', '-filetype=obj', '-mtgpu-enable-const-calc=1', '-o', '/tmp/triton-musa-8y4br0gb/kernel.o']

['/usr/local/lib/python3.10/dist-packages/triton/backends/mthreads/bin/ld.lld', '-flavor', 'gnu', '-shared', '/tmp/triton-musa-8y4br0gb/kernel.o', '-o', '/tmp/triton-musa-8y4br0gb/kernel.mubin']

MTCC is not available inside my Docker container and llc version is:

LLVM (http://llvm.org/):
  LLVM version 14.0.0
  Optimized build.
  Default target: x86_64-unknown-linux-gnu
  Host CPU: znver3

  Registered Targets:
    mtgpu  - The MT MTGPU GPUs
    x86    - 32-bit X86: Pentium-Pro and above
    x86-64 - 64-bit X86: EM64T and AMD64

May I ask you two questions:

1. Is this llc equivalent to MTCC?
2. Is there a feasible way to verify whether the load <4 x float> vectorized instruction is actually compiled into the final .mubin binary file?

[zhzhcookie]

Please try compiling flagtree mthreads3.6 using a newer version after #839.