[Feature Request] Enhance TileLang to auto inject ifthenelse for Out-of-Bounds Memory Access in TileLang

[LeiWang1999]

for example, we can handle a small m with large BlockM with the following solution:

for i, j in T.Parallel(block_M, block_N):
    m, n = by * block_M + i, bx * block_N + j
    if m < M and n < N:
        C[m, n] = C_shared[
            i // micro_size_x,
            j // micro_size_y,
            i % micro_size_x,
            j % micro_size_y,
        ]

It will be great to allow T.Parallel auto analysis the memory access and inject the if then else node by itself.

[LeiWang1999]

To address the need for auto-injection of if-then-else checks for out-of-bounds memory access in TileLang, we’ve implemented a new pass: tl.LegalizeSafeMemoryAccess. This pass is designed to analyze memory accesses within loops, identify potential out-of-bounds conditions, and automatically insert the appropriate if-then-else statements to guard such accesses.

Key Features of the Pass:

1. Global Memory Access Detection: The pass scans for BufferLoad and BufferStore operations targeting global memory and analyzes their access patterns.
2. Bounds Checking: It compares indices against the buffer’s shape dimensions and automatically generates conditions to ensure access safety.
3. Condition Injection: If the pass detects that a memory access could potentially be out of bounds, it wraps the loop body in the necessary if-then-else conditions to enforce valid access.
4. Leaf Loop Analysis: The pass only applies the transformations to loops that do not have inner loops, ensuring that it does not over-constrain the program unnecessarily.
5. Simplification-Ready Conditions: The generated conditions are compatible with downstream simplification passes, ensuring they are optimized efficiently.

Here’s an example of how the pass transforms a loop:

Input:

for i, j in T.Parallel(block_M, block_N):
    m, n = by * block_M + i, bx * block_N + j
    C[m, n] = C_shared[
        i // micro_size_x,
        j // micro_size_y,
        i % micro_size_x,
        j % micro_size_y,
    ]

Output (after the pass):

for i, j in T.Parallel(block_M, block_N):
    m, n = by * block_M + i, bx * block_N + j
    if m < M and n < N:
        C[m, n] = C_shared[
            i // micro_size_x,
            j // micro_size_y,
            i % micro_size_x,
            j % micro_size_y,
        ]

The injected condition ensures safe memory access by checking that m < M and n < N before accessing C[m, n].

How to Use the Pass

The pass can be applied as part of the TVM compilation pipeline by invoking tl.LegalizeSafeMemoryAccess. The registration of the pass allows it to be seamlessly integrated into existing workflows. For example:

    mod = tl.transform.LegalizeSafeMemoryAccess()(mod)
    # Inject Simplify to remove the duplicated conditions
    # which may be introduced by the LegalizeSafeMemoryAccess
    mod = tir.transform.Simplify()(mod)

This should address scenarios where loops access memory blocks with dimensions larger than the actual data, such as in the case of padding or tiling.

[LeiWang1999]

closed as has been implemented by pr #267