[Benchmark] gather_gemv kernel and test

[Sibylau]

Stacked PRs:

• Run assertExpectedJournal for gather_gemv when is_cuda #656
• ->[Benchmark] gather_gemv kernel and test #635

---

[Benchmark] gather_gemv kernel and test

[Sibylau]

• code generation test passes with python -m unittest test_examples.TestExamples.test_gather_gemv

• tritonbench accuracy check passes with python benchmarks/run.py --metrics accuracy --kernel gather_gemv (after applying @yf225 's patch for test_eager: gather_gemv: fix eager impl to support large input; fix compile impl meta-pytorch/tritonbench#455 )

[yf225]

thanks @Sibylau ! left some nit comments, and also might need to rebase to fix conflicts with main branch

[yf225]

maybe can remove this comment

[yf225]

maybe remove?

[yf225]

Not exactly sure why AMD job fails.. but I suspect changing the block sizes to some smaller value might help

[Sibylau]

@yf225 The CI test on rocm fails due to code mismatch:

-         dot = tl.dot(tl.cast(gathered, tl.float32), tl.cast(load_1, tl.float32), input_precision='tf32', out_dtype=tl.float32)
+         dot = tl.split(tl.permute(tl.reshape(tl.split(tl.permute(tl.reshape(tl.split(tl.permute(tl.reshape(tl.split(tl.permute(tl.reshape(tl.split(tl.permute(tl.reshape(tl.dot(tl.reshape(tl.permute(tl.join(tl.cast(gathered, tl.float32), 

Do you know why the generated code for AMD is different? can i put a @skipIfRocm for this kernel test?

[yf225]

Instead of using skipIfRocm which skips the whole test including the output equality check, maybe we can add a skip_rocm: bool arg to def assertExpectedJournal that skips the journal check if the device is rocm. And to check the device is rocm, we can add something like this to _testing.py:

def is_rocm() -> bool:
    """Return True if running on ROCm (AMD GPU)."""
    return (
        triton.runtime.driver.active.get_current_target().backend == "hip"
        and DEVICE.type == "cuda"
    )

(Please feel free to do this in a follow-up PR. Thanks!)