[inductor] Fix a use before def error

[desertfire]

Summary: for #918

[desertfire]

BEFORE: tmp4 = tl.load(in_ptr1 + 0, tmp2) is incorrectly promoted out of the reduction loop.

@reduction_heuristics(size_hints=[32, 8192])
@triton.jit
def kernel0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK, 1])
    xmask = xindex < xnumel
    rbase = tl.reshape(tl.arange(0, RBLOCK), [1, RBLOCK])
    x0 = xindex
    tmp4 = tl.load(in_ptr1 + 0, tmp2)
    _tmp12 = tl.zeros([XBLOCK, RBLOCK], tl.float32) + 0
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp0 = r1 + (7823*x0)
        tmp1 = 140800
        tmp2 = tmp0 < tmp1
        tmp3 = tl.load(in_ptr0 + (r1 + (7823*x0)) % 140800 + tl.zeros([XBLOCK, RBLOCK], tl.int32), xmask & rmask & tmp2, eviction_policy='evict_last')
        tmp5 = tmp3 / tmp4
        tmp6 = tl.load(in_ptr2 + (r1 + (7823*x0)) % 140800 + tl.zeros([XBLOCK, RBLOCK], tl.int32), xmask & rmask & tmp2, eviction_policy='evict_last')
        tmp7 = tmp5 * tmp6
        tmp8 = 0
        tmp9 = tmp2 | tl.zeros(tmp7.shape, tmp2.dtype) if tmp7.numel > 1 else tmp2
        tmp10 = tmp9 | tl.zeros(tmp8.shape, tmp9.dtype) if tmp8.numel > 1 else tmp9
        tmp11 = tl.where(tmp10, tmp7, tmp8)
        _tmp12 = tl.where(xmask & rmask, _tmp12 + tmp11, _tmp12)
    tmp12 = tl.reshape(tl.sum(_tmp12, 1), [XBLOCK, 1])
    tl.store(out_ptr0 + x0, tmp12, xmask)

AFTER:

@reduction_heuristics(size_hints=[32, 8192])
@triton.jit
def kernel0(in_ptr0, in_ptr1, in_ptr2, out_ptr0, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK, 1])
    xmask = xindex < xnumel
    rbase = tl.reshape(tl.arange(0, RBLOCK), [1, RBLOCK])
    x0 = xindex
    _tmp12 = tl.zeros([XBLOCK, RBLOCK], tl.float32) + 0
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp0 = r1 + (7823*x0)
        tmp1 = 140800
        tmp2 = tmp0 < tmp1
        tmp3 = tl.load(in_ptr0 + (r1 + (7823*x0)) % 140800 + tl.zeros([XBLOCK, RBLOCK], tl.int32), xmask & rmask & tmp2, eviction_policy='evict_last')
        tmp4 = tl.load(in_ptr1 + 0, tmp2)
        tmp5 = tmp3 / tmp4
        tmp6 = tl.load(in_ptr2 + (r1 + (7823*x0)) % 140800 + tl.zeros([XBLOCK, RBLOCK], tl.int32), xmask & rmask & tmp2, eviction_policy='evict_last')
        tmp7 = tmp5 * tmp6
        tmp8 = 0
        tmp9 = tmp2 | tl.zeros(tmp7.shape, tmp2.dtype) if tmp7.numel > 1 else tmp2
        tmp10 = tmp9 | tl.zeros(tmp8.shape, tmp9.dtype) if tmp8.numel > 1 else tmp9
        tmp11 = tl.where(tmp10, tmp7, tmp8)
        _tmp12 = tl.where(xmask & rmask, _tmp12 + tmp11, _tmp12)
    tmp12 = tl.reshape(tl.sum(_tmp12, 1), [XBLOCK, 1])
    tl.store(out_ptr0 + x0, tmp12, xmask)

[desertfire]

There is still another error where tl.load fails for a dim=0 tensor,

tmp4 =  tl.load(in_ptr1 + 0, tmp2)
...
arg1_1 = rand_strided((), (), device='cuda', dtype=torch.float32)

[jansel]

The index != 0 here seems related, I think this was related to loading the random mask.

[jansel]

Why not

Suggested change

	if (need_dense and not have_dense) or load_index0:
	if (need_dense and not have_dense) or index == 0:

?

[desertfire]

I wasn't sure if this is safe to do if we are not calling from load, but CI run suggests this is ok.

[ngimel]

There is still another error where tl.load fails for a dim=0 tensor,

wdym? How does it fail?

[desertfire]

There is still another error where tl.load fails for a dim=0 tensor,

wdym? How does it fail?

The error I saw was,

>       name, asm, shared_mem = _triton.code_gen.compile_ttir(backend, generator.module, device, num_warps, num_stages, extern_libs)
E       MemoryError: std::bad_alloc

../triton/python/triton/code_gen.py:1320: MemoryError

Adding tl.zeros([XBLOCK, RBLOCK] to the indexing would make it pass, but as @ngimel pointed out in #981, what we really need here is mask as None.