[TKW] Detect contiguous access pattern in mapped reads/writes

[Hardcode84]

Add utility to detect contiguous access pattern in mapped reads/writes so we can use contiguous vector ops instead of gathers/scatters during lowering.

Iterate over elements_per_thread and check if every subsequent mapped index have only diff of 1 in fastest changing dim.

Also, it need some custom sympy simplifications to successfully work on IGEMM conv.

[harsh-nod]

How about something like this?

[harsh-nod]

To accomplish this, you would do the following

> next_indices[2] - prev_indices[2]
-floor(floor((Mod($T0, 4))/80 + floor($ARGC*HF*WF*BLOCK_K/8)/80)/3) + floor(floor((Mod($T0, 4))/80 + floor($ARGC*HF*WF*BLOCK_K/8)/80 + 1/640)/3)

> sympy.cse(next_indices[2] - prev_indices[2])[0]
[(x0, (Mod($T0, 4))/80 + floor($ARGC*HF*WF*BLOCK_K/8)/80)]
> sympy.cse(next_indices[2] - prev_indices[2])[1]
[-floor(floor(x0)/3) + floor(floor(x0 + 1/640)/3)]

Then, finally if you want to add the assumption that x0 ~ O(1), you can just subs x0 -> 1 and you will get 0.

[Hardcode84]

So, as any sane person should do, I wrote a fuzzing tests (see latest commit) and it indeed found some failing cases. But after stricten some checks (and marking induction vars as non-negative) and running 100 000s tests I'm now reasonable sure my transformations are correct.

[harsh-nod]

Thanks that sounds good, but I think we have been thinking about this the wrong way.

So essentially what we are looking to do here is determine whether the mapping is "contiguity-preserving" (which means that elements that are contiguous pre-mapping will remain contiguous post-mapping). This is a function entirely of the mapping and should not depend at all on the particular access pattern of the underlying data. So I think these patterns to simplify the index sequences are not required and instead we should be looking just at the mapping. For example, looking at one IGEMM mapping which maps a tensor of shape [M, K] to a tensor of shape [N, H, W, C]

x_mapping = tkw.IndexMapping(
        num_iterators=2,
        inputs={
            N: i // SZ_OUT,
            C: j % C,
            H: (i % SZ_OUT) % W_OUT * stride + (j // C) % WF,
            W: (i % SZ_OUT) // W_OUT * stride + (j // C) // WF,
        },
        outputs={M: i, K: j},
    )

we can immediately deduce that this mapping is contiguity-preserving because j gets mapped to j % C.

[Hardcode84]

yeah, sympy.cse looks useful, but I weren't able to construct working code quickly. I will need to think more on it offline.

[harsh-nod]

sounds good, let's land this PR. I have a WIP PR that should simplify this a lot use sympy.cse.

[harsh-nod]

Just some minor changes, but overall looks good. How did the cse simplification help?

[harsh-nod]

let's land for now. I think we can greatly simplify the checks by using sympy.cse but we can explore that after landing this PR.