Add Layernorm kernel

[rahulbatra85]

No description provided.

[brunomazzottiamd]

@rahulbatra85, I can't see anything wrong with your PR. I've just have some questions and minor code cleanup suggestions. Feel free to ignore them if you judge appropriate.

Please drop a short line about this new kernel to python/perf-kernels/README.md file.

[brunomazzottiamd]

What do you think about running all tests with just one pytest invocation? According to https://docs.pytest.org/en/stable/how-to/usage.html, it's possible to do something like pytest -vvvv ./python/perf-kernels. By this way, we'll be editing .github/workflows/amd_perf_kernel_Integration_tests.yml less often and new tests are going to run by default. Do you see any drawback?

Maybe it's worth asking @micmelesse's opinion on this.

[rahulbatra85]

yeah that's a Michael question

[brunomazzottiamd]

Let's wait for Michael's opinion!

[micmelesse]

that is fine. I think some of the tests are broken but maybe worth it to see the state of things

[brunomazzottiamd]

Just an idea:
We have three for loops that do masked loads. Do you foresee any benefit of peeling the last iteration of each loop so all iterations except the last one do unmasked loads? I think Shucai and Xiaohu got some performance improvements doing this with GEMMs. I'm not sure if the idea could be beneficial for layer norm.

[rahulbatra85]

ok, yeah, I didn't think of that. Will try this out

[brunomazzottiamd]

Please let me know if this helped at all.

[scxiao]

Can you add an input use_mask: tl.constexpr to the kernel? Then in the implementation, the read from global memory can be like:

    loop_num = tl.cdiv(n_cols, BLOCK_SIZE)
    if use_mask:
        loop_num -= 1
    #calculate mean
    mean = 0
    _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for b in range(0, loop_num):
        col_offsets = b * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        x_block = tl.load(x_ptr_start + col_offsets).to(tl.float32)
        _mean += x_block
    if use_mask:
        col_offsets = loop_num * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        x_block = tl.load(x_ptr_start + col_offsets, mask=col_offsets < n_cols, other=0.).to(tl.float32)
        _mean += x_block
    mean = tl.sum(_mean, axis=0) / n_cols

    #same for the variance calculation.
    _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    for b in range(0, n_cols, BLOCK_SIZE):
        col_offsets = b + tl.arange(0, BLOCK_SIZE)
        x_block = tl.load(x_ptr_start + col_offsets, mask=col_offsets < n_cols, other=0.).to(tl.float32)
        x_block = tl.where(col_offsets < n_cols, x_block - mean, 0.)
        _var += x_block * x_block
    var = tl.sum(_var, axis=0) / n_cols
    rstd = tl.rsqrt(var + eps)

In this way, we do need mask in most of the iterations, which can make the load input be more efficient.

[rahulbatra85]

ok, will try this out.

[rahulbatra85]

I did this and the perf improves https://github.com/ROCm/triton-internal/issues/126#issuecomment-2369175077