[FRONTEND] Fix arg name conflict bug

[Sarbojit2019]

There is a bug introduced in latest triton code where AttributeError is getting thrown if kernel function args has same name as global variable. See the below example for more clarity

import torch
import triton
import triton.language as tl

@triton.jit
def kernel(
      out, adj, num_rows,
      WG_SIZE : tl.constexpr
  ):
    zero_indices = tl.zeros((WG_SIZE,), dtype=adj.dtype.element_ty)


if __name__ == '__main__':
    device = torch.device('cuda')

    num_rows, num_indices, heads = 16, 16, 16, 4
    adj = torch.zeros(num_rows + 1, dtype=torch.int32, device=device)
    out = torch.empty((num_indices,heads), dtype=e_src.dtype, device=e_src.device)
    # invoke triton kernel
    launch_grid = (num_rows * heads // 256,)
    kernel[launch_grid](out, adj, num_rows, 256) 

Error it throws something like below
return getattr(lhs, node.attr) ^^^^^^^^^^^^^^^^^^^^^^^ AttributeError: 'torch.dtype' object has no attribute 'element_ty'

With the proposed change I see test is passing. I have verified it on Nvidia and Intel platforms.

[yiakwy-xpu-ml-framework-team]

@Sarbojit2019 I found python AST generator has bug in estimating SRAM used and tensor assignment, did you noticed that before ?

[Jokeren]

Taking a look now (the example itself seems buggy).

[Jokeren]

@Sarbojit2019 Please correct your reproducer so that we can use it without any changes

[Sarbojit2019]

@Jokeren, could you please help me understand the issue you see with the reproducer?

[Jokeren]

We cannot run your reproducer directly. There are at least two problems IIRC. For example, e_src is not defined.

[Jokeren]

Take a quick look, another one should be this:

    num_rows, num_indices, heads = 16, 16, 16, 4

[Jokeren]

I think it is incomplete because we can also define a local variable with assignment, not as function arguments.

For example:

@triton.jit
def kernel(...):
    ptr = tl.load(...)
    print(ptr.dtype.element_ty)

ptr = torch.tensor(...)

[Jokeren]

@Sarbojit2019 Do you plan to address the issue in this PR?

[Sarbojit2019]

Yes, proposed the fix in latest commit.

[Jokeren]

Thanks for proposing the fix.

This is a known problem previously reported by PT2 developers. We were sticking to the "implicit rule" that names cannot be same in the local and the global scopes. But I think it's a good try to fix it now.

[Sarbojit2019]

@Jokeren, I have updated the sample that I had pasted earlier.

import torch
import triton
import triton.language as tl

@triton.jit
def kernel(
      out, adj, num_rows,
      WG_SIZE : tl.constexpr
  ):
    zero_indices = tl.zeros((WG_SIZE,), dtype=adj.dtype.element_ty)


if __name__ == '__main__':
    device = torch.device('cuda')

    num_rows, num_cols, num_indices, heads = 16, 16, 16, 4
    e_src = torch.randn((num_cols,heads), dtype=torch.float32, device=device) + 0.5
    adj = torch.zeros(num_rows + 1, dtype=torch.int32, device=device)
    out = torch.empty((num_indices,heads), dtype=e_src.dtype, device=e_src.device)
    # invoke triton kernel
    launch_grid = (num_rows * heads // 256,)
    kernel[launch_grid](out, adj, num_rows, 256)

[Jokeren]

Hi @Sarbojit2019 , I had a deeper thought.

The naming conflict won't be fully resolved until all possible definitions in AST visitor has been covered. For example, the iteration variable of a for loop could also shadow a global variable.

So I think we need to have a thorough test under test/runtime to cover all scenarios. It could be done in the next PR. Would you like to contribute to this feature?

[Sarbojit2019]

@Jokeren, I would be happy to work on it.
As I understood triton supports smaller set of operations compare to python e.g. simultaneous assignments is not supported in triton. Could you please let me know where can I find the grammar for triton? This will help fix possible scenarios better.

[Jokeren]

It would be a simpler version of CodeGenerator, you only need to care about where new variables are defined.