commit 337c8aaf544053ac512bb7c023c572f8e2d51e7e
author Max191 <44243577+Max191@users.noreply.github.com> Wed Aug 20 08:10:09 2025 -0700
committer GitHub <noreply@github.com> Wed Aug 20 08:10:09 2025 -0700
tree 3839cfd34dadcee5c37e09d4b3126259782d838d
parent 1f7762d884d26cbc2faf435bb393651c559518f1

[Codegen] Improve early bufferized padding codegen (#21694)

The current way we handle padding in CombineLayoutTransformation (the
early bufferization path) turns out to be very inefficient in some
cases, because we spend a lot of time spinning on indexing computation
while iterating over the output buffer. This PR updates the way we
generate the IR for writing padding values to avoid iterating over the
non-padded parts of the output buffer. The old and new methods are
described below:

### Old method ###

For a given tensor with padding, we would simply iterate over the entire
tensor, and write padding values when the indices fell within the padded
part of the output buffer. Visually this looks like the following:
```
          1st (only) loop (`x` == masked off)
         |———————————————————————————————————| 
                            |
                            v
         +———————————————————————————————————+
         | . . . . . . . . . . . . . . . . . |
         | . +———————————————————————+ . . . |
         | . | x x x x x x x x x x x | . . . |
         | . | x x x x x x x x x x x | . . . |
         | . | x x x x x x x x x x x | . . . |
         | . | x x x x x x x x x x x | . . . |
         | . +———————————————————————+ . . . |
         | . . . . . . . . . . . . . . . . . |
         +———————————————————————————————————+
```
### New method ###

There are multiple distributed loops, each loop performing the padding
for a hyperplane spanning the low or high pad of a corresponding
dimension, and fully spanning all other dimensions of the tensor. This
means that there is some overlap on the intersection of high and low
padding for each dimension.

```
 1st loop (low pad d0)        2nd loop (high pad d0)
        |———|                       |———————|
          |                             |
          v                             v
        +———+———————————————————————+———————+     ———
        | . | . . . . . . . . . . . | . . . | <--  |  3rd loop (low pad d1)
        +———+———————————————————————+———————+     ———
        | . | . . . . . . . . . . . | . . . |
        | . | . . . . . . . . . . . | . . . |
        | . | . . . . . . . . . . . | . . . |
        | . | . . . . . . . . . . . | . . . |
        +———+———————————————————————+———————+     ———
        | . | . . . . . . . . . . . | . . . | <--  |  4th loop (high pad d1)
        +———+———————————————————————+———————+     ———
```

With the new method, there is some overlap between loops, but the
overall padding computation is much cheaper because padding is expected
to be small relative to the output buffer size.

---------

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>