tree d2f96f292054b8ce93c3ca19ec15bade05718541
parent 78005ef66ecb75f984d1d24e46bcaf6d01ed7317
author Han-Chung Wang <hanhan0912@gmail.com> 1713572668 -0700
committer GitHub <noreply@github.com> 1713572668 -0700
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJmIws8CRC1aQ7uu5UhlAAAqaYQAIUtr0qXGD5x0uUHYpP1IWLx
 viMsYpmehc6wI+8r4oVYrck0h+K2FoPndsdLbeOukszG58hTeMRXGPeFyN6X7/oP
 Qdf8SGRKpgCbP4gihRAsR6v+9+CoYibnT1Q5j1YOZk+oQd1Ms3NkWJlRFtWCAI42
 /+V+cuPPDh7mlcy9ljFx+AGDw/vYdIfLpuM6mMLtQ36G36e3AJndmihoOJm1cxtw
 O/RF6aWm/rRjcCYCDeY/nleImPFIxikr9bL1dR02no3Ay+gcqGyjqneCSFzQsfHe
 n6b740qbSvNH5uqb1ZiuNvRQwH8bKESkDU4hybLPEeppGhhVx7lJsSECJp1zFtw2
 QbRamR13eYEFpy5QbKsMBgYOpaVq26yQMr8HNeTOfwFfsGtn3iu5dgQC3v5nKMRE
 W6ype9z758Egmu1nKnKChj4tqJp6W67uGCm7uNXooFRAZBzjJrl+T9i/V5pjvycr
 IF4woxTAFb/ZPSfGfnqntURTev76FovXAUi2D09OMg1FNAApQVEZTb0FjhEAsKsh
 J3cr76nkr6UuPk6qBKZWouycCJTVlfRJ9c066y6sVfGqWAj6wODcCmQDsvWLiTD9
 xphkKQ5gXRhNbCrOQ48vDc2CAzybvcU9RTHa29+IVmarg6AtpcKMQeUT/KpujWkt
 yej79tDLCC+xocIRPSMU
 =BUA1
 -----END PGP SIGNATURE-----
 

[DT] Teach encoding about padding. (#17077)

The revision has four major commits. They are expected to be landed
together because we need all the piece to make it work.

## Make encodings be able to carry padding semantics

https://github.com/openxla/iree/pull/17077/commits/7967faf4d2b442f7f1229d3801dbc95678ebe051
introduces a `round_dims_to` integer array on encodings. It represents
the values for padding M,N,K dimensions. This provides the hints for
both host and device the values that every
dimension should be aligned with. Eventually we should have a better way
for propagating the information between host side and device side. The
revision is a step towards to the goal.

The commit adds an option to SetEncoding pass. If the `padFactor` is
set, the `round_dims_to` will be filled with the values; it only
generates set_encoding ops, but not
`iree_linalg_ext.upper_bound_tile_size` and `tensor.pad` ops.

## Teach Pack/UnPack Materialization Patterns about the new field

https://github.com/openxla/iree/pull/17077/commits/2652c028b586b984dca237cfc6cd53a1ffa5235e
teaches the materialization patterns to handle the new field. If the
field is set, the inner tile sizes can't be greater than corresponding
`round_dims_to` values. Otherwise, the actual buffer size could
mismatch.

## Teach stream.tensor.sizeof to take encoding into accounts

https://github.com/openxla/iree/pull/17077/commits/365dc413250675ae8e49cfcaccc3e4341d1d3432
teaches stream.tensor.sizeof to calculate proper sizes based on
encodings. The encodings have `role`, `indexing_maps`, and
`round_dims_to`. So it is able to look at `role` and `indexing_map` to
infer contraction dimensions; pads the dimension to be aligned with
values in `round_dims_to`. E.g.,

```mlir
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
util.func public @sizeof_lhs_encoding_dynamic(%arg0: index, %arg1: index) -> index {
  %0 = stream.tensor.sizeof tensor<?x?xf32, #iree_linalg_ext.encoding<
    role = LHS,
    element_types = [f32, f32, f32],
    original_type = tensor<?x?xf32>,
    user_indexing_maps = [#map, #map1, #map2],
    round_dims_to = 4, 8, 16>>{%arg0, %arg1} : index
  util.return %0 : index
}
// CHECK-LABEL: @sizeof_lhs_encoding_dynamic
// CHECK-DAG:     %[[C4:.+]] = arith.constant 4 : index
// CHECK-DAG:     %[[C16:.+]] = arith.constant 16 : index
// CHECK:         %[[CEIL_DIV_D0:.+]] = arith.ceildivui %arg0, %[[C4]]
// CHECK:         %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C4]]
// CHECK:         %[[CEIL_DIV_D1:.+]] = arith.ceildivui %arg1, %[[C16]]
// CHECK:         %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
// CHECK:         %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
// CHECK:         %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
// CHECK:         return %[[T1]]
```

## Add e2e tests

https://github.com/openxla/iree/pull/17077/commits/5ca4a4b7a5a672e25d257cc9443c030735c305e5
adds a new test suite with
`--iree-global-opt-enable-early-materialization=false`, so we have
enough e2e test coverage for the new path.