[Flow] Make sink reshapes changes less conservative. #17706
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
MaheshRavishankar
force-pushed
the
sdxl_quantized_sink_reshape_fixes
branch
from
3a50abd
to
e3066d6
Compare
MaheshRavishankar
added
benchmarks:cuda
Abbreviated Benchmark Summary@ commit 319360d130b43a73ace1fba52c56f4b90db2b5f0 (vs. base 7090f64b7bd60a597ee6c61b8ff0d624153cb7f0) Data-Tiling Comparison TableClick to show
Regressed Latencies 🚩
Improved Latencies 🎉
[Top 3 out of 13 results showed] Regressed Stream IR Dispatch Count (# of cmd.dispatch ops) 🚩
For more information: |
MaheshRavishankar
force-pushed
the
sdxl_quantized_sink_reshape_fixes
branch
2 times, most recently
from
6cae1e9
to
5c9dac6
Compare
qedawkins
approved these changes
Jun 26, 2024
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
|
Have been trying to repro it. Can't repro locally. |
|
I think the regression can be seen here: Afterfull after change %expanded_181 = tensor.expand_shape %43 [[0, 1], [2]] output_shape [1, 4, 768] : tensor<4x768xf32> into tensor<1x4x768xf32>
%expanded_182 = tensor.expand_shape %46 [[0, 1]] output_shape [1, 4] : tensor<4xf32> into tensor<1x4xf32>
%expanded_183 = tensor.expand_shape %45 [[0, 1]] output_shape [1, 4] : tensor<4xf32> into tensor<1x4xf32>
%47 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]}
ins(%expanded_181, %expanded_182, %cst_60, %cst_59, %expanded_183
: tensor<1x4x768xf32>, tensor<1x4xf32>, tensor<768xf32>, tensor<768xf32>, tensor<1x4xf32>)
outs(%7 : tensor<1x4x768xf32>) {
^bb0(%in: f32, %in_584: f32, %in_585: f32, %in_586: f32, %in_587: f32, %out: f32):
%355 = arith.divf %in_584, %cst_9 : f32
%356 = arith.addf %355, %cst_8 : f32
%357 = math.rsqrt %356 : f32
%358 = arith.mulf %357, %in_585 : f32
%359 = arith.mulf %in_587, %358 : f32
%360 = arith.subf %in_586, %359 : f32
%361 = arith.mulf %in, %358 : f32
%362 = arith.addf %361, %360 : f32
linalg.yield %362 : f32
} -> tensor<1x4x768xf32>
%collapsed_184 = tensor.collapse_shape %47 [[0, 1], [2]] : tensor<1x4x768xf32> into tensor<4x768xf32>It should be possible to sink the expanded dims and fold with the collapse shape, but since each input indexing map is a projection the |
|
Thanks Ian... let me look into this a little bit more. |
MaheshRavishankar
force-pushed
the
sdxl_quantized_sink_reshape_fixes
branch
2 times, most recently
from
41ed0e1
to
9bb15c6
Compare
Max191
reviewed
Jun 28, 2024
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
compiler/src/iree/compiler/Dialect/Flow/Transforms/SinkReshapes.cpp
Outdated
Show resolved
Hide resolved
While deciding if a reshape needs "sinking", for a `tensor.expand_shape` -> `linalg.*`, first check was to check that the `linalg.*` operation could already fuse with one of its existing producers. That check was broadly aggressive. The fusion only kicks in when the iteration domains match. Eventually the actual dispatch formation logic needs to be commoned to a single place to do this better, but kicking that to a follow up. Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
MaheshRavishankar
force-pushed
the
sdxl_quantized_sink_reshape_fixes
branch
from
9bb15c6
to
253aecc
Compare
Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
MaheshRavishankar
force-pushed
the
sdxl_quantized_sink_reshape_fixes
branch
from
b808b9b
to
f8163b2
Compare
Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com>
|
Ok, folks. I am landing this with the one regression in number of dispatches. I am checking some out of tree models and this seems to do better. |
LLITCHEV
pushed a commit
to LLITCHEV/iree
that referenced
this pull request
Jul 30, 2024
While deciding if a reshape needs "sinking", for a `tensor.expand_shape` -> `linalg.*`, first check was to check that the `linalg.*` operation could already fuse with one of its existing producers. That check was broadly aggressive. The fusion only kicks in when the iteration domains match. Eventually the actual dispatch formation logic needs to be commoned to a single place to do this better, but kicking that to a follow up. --------- Signed-off-by: MaheshRavishankar <mahesh.ravishankar@gmail.com> Signed-off-by: Lubo Litchev <lubol@google.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
While deciding if a reshape needs "sinking", for a
tensor.expand_shape->linalg.*, first check was to check that thelinalg.*operation could already fuse with one of its existing producers. That check was broadly aggressive. The fusion only kicks in when the iteration domains match. Eventually the actual dispatch formation logic needs to be commoned to a single place to do this better, but kicking that to a follow up.