[Mlir-commits] [mlir] [MLIR][NFC] Add allow Insert/extract slice option to pack/unpack op (PR #117340)

[Andrzej Warzyński]

banach-space wrote:

> I see discussions converging here. My next commit is going to be:
> 
> * Get rid of the unit test

I would keep the unit test. 

When I asked about the test that you added (and subsequently removed), it just wasn't clear to me what to look for and what the difference with the existing tests should be. The comment from @Max191 makes that very clear:
```mlir
// Current
func.func @pack_as_pad_default() {
  %pad = tensor.pad %src ...
  %producer = tensor.insert_slice %pad into %dest
}

// New
func.func @pack_as_pad_{disable|no|skip}_insert_slice() {
  %pad = tensor.pad %src
  %expand = tensor.expand_shape %pad ...
  %producer = linalg.transpose ins(%expand ..
}
```

**Suggestion 1**
* This PR should include a test that follows the pattern above ^^^. 

I think that the following test would be sufficient and very illustrative (trimmed for brevity - feel free to re-use/expand):

```mlir
// CHECK-LABEL: func.func @pack_disallowed_as_pad(
func.func @pack_disallowed_as_pad(%arg0, %arg1) {

  //      CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] 
  //  CHECK-NOT: %[[RES:.*]] = tensor.insert_slice %[[PAD]]
  //      CHECK: %[[PAD_EXPANDED:.*]] = tensor.expand_shape %[[PAD]]
  //      CHECK: %[[RES:.*]] = linalg.transpose %[[PAD_EXPANDED]]
  %pack = tensor.pack %arg0 padding_value(%cst_0 : f32) inner_dims_pos = [0, 1, 2, 3] inner_tiles = [136, 64, 16, 16] into %arg1
    : tensor<129x47x16x16xf32> -> tensor<1x1x1x1x136x64x16x16xf32>
  return %pack :  tensor<1x1x1x1x136x64x16x16xf32>
}
```
--- 

IIUC, these are the existing tests that exercise similar path: https://github.com/llvm/llvm-project/blob/1b2c8f104f9c6f26500ab608060bbc6b7f40f5e1/mlir/test/Dialect/Linalg/transform-lower-pack.mlir#L3-L60

Both of these tests are called `func.func @pack`, so it's hard to tell what the difference is? Are they duplicating one another? IIUC, these are not "[pad-like](https://github.com/llvm/llvm-project/blob/1b2c8f104f9c6f26500ab608060bbc6b7f40f5e1/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp#L298-L327)"? 

**Suggestion 2**
* Update the name of these tests ^^^ with something more descriptive so that it's easier to identify the edge cases being tested.


---

> I'm not really sure where a test like this belongs. If you have a suggestion for where we could put it that would be very helpful!

Sure, happy to help! I see multiple options:

**Option 1:** Add a test file in test/Dialect/Linalg (e.g. "transform-lower-pack-and-fuse.mlir") and test it there. Looks like there's already [transform-tile-and-fuse.mlir](https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/transform-tile-and-fuse.mlir), so use that as your source of inspiration :)

**Option 2:** Some e2e tests exercise complex compilation pipelines, e.g. [pack-dynamic-inner-tile.mlir](https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/pack-dynamic-inner-tile.mlir). But that's quite a heavy hammer.

**Option 3:** Add a bit more powerful TD Op that would include fusion and various related patterns that create fusion opportunities.

I'd start with **Option 1** and leave **Option 2** and **Option 3** for the future. On a related note, we should start creating sub-directories in test/Dialect/Linalg/ - there's too many tests prefixed "transform-{}.mlir" 😅 

---

Hope this helps and thanks for working on this 🙏🏻 !

https://github.com/llvm/llvm-project/pull/117340