Skip to content

[GlobalOpt] Keep strided input maps when the producer can fuse - #24950

Open
zhangpan2001 wants to merge 1 commit into
iree-org:mainfrom
zhangpan2001:fix/24752-dequant-strided-conv-fusion
Open

zhangpan2001 wants to merge 1 commit into
iree-org:mainfrom
zhangpan2001:fix/24752-dequant-strided-conv-fusion

Conversation

@zhangpan2001

Copy link
Copy Markdown
Contributor

Factoring dim * cst strides out of a contraction-like generic inserts a strided tensor.extract_slice between the contraction and its input. When the input has an elementwise producer (a dequantize feeding a 1x1 stride-2 convolution, say), the slice blocks both fusion paths: elementwise fusion requires the producer to directly define the consumer's operand, and tile-based producer fusion only supports unit-stride slices. Codegen falls back to materializing the producer as a whole-tensor stack allocation in every workgroup: a 3.2 MB alloca for a 1x256x56x56xf32 dequantize, well over the 32 KB stack limit, and the compile fails.

Skip the factoring only when the strided input's producer can fuse elementwise into the contraction and into every other user of its result, so the fusion actually eliminates the producer. Codegen then merges the producer into the contraction, which reads the input directly through the strided map with no intermediate tensor. Producers that are materialized anyway (e.g. a result also returned from the function) keep the factored projected-permutation input map.

Strided contractions with a fusable producer now lower through the generic path instead of contraction selection such as AMDGPU MFMA, the original motivation for the factoring, in exchange for eliminating the intermediate tensor. Strided contractions on dispatch-boundary inputs keep the factored form.

Fixes #24752

Tests: iree-opt --split-input-file --mlir-print-local-scope -iree-global-opt-convert-strided-contraction-to-contraction compiler/src/iree/compiler/GlobalOptimization/test/strided_contraction_to_contraction.mlir

Factoring `dim * cst` strides out of a contraction-like generic inserts a
strided `tensor.extract_slice` between the contraction and its input. When
the input has an elementwise producer (a dequantize feeding a 1x1 stride-2
convolution, say), the slice blocks both fusion paths: elementwise fusion
requires the producer to directly define the consumer's operand, and
tile-based producer fusion only supports unit-stride slices. Codegen falls
back to materializing the producer as a whole-tensor stack allocation in
every workgroup: a 3.2 MB alloca for a 1x256x56x56xf32 dequantize, well
over the 32 KB stack limit, and the compile fails.

Skip the factoring only when the strided input's producer can fuse
elementwise into the contraction and into every other user of its result,
so the fusion actually eliminates the producer. Codegen then merges the
producer into the contraction, which reads the input directly through the
strided map with no intermediate tensor. Producers that are materialized
anyway (e.g. a result also returned from the function) keep the factored
projected-permutation input map.

Strided contractions with a fusable producer now lower through the generic
path instead of contraction selection such as AMDGPU MFMA, the original
motivation for the factoring, in exchange for eliminating the intermediate
tensor. Strided contractions on dispatch-boundary inputs keep the factored
form.

Fixes iree-org#24752

Tests:
- iree-opt --split-input-file --mlir-print-local-scope
  -iree-global-opt-convert-strided-contraction-to-contraction
  compiler/src/iree/compiler/GlobalOptimization/test/strided_contraction_to_contraction.mlir

Co-authored-by: OpenAI <noreply@openai.com>
Signed-off-by: zhangpan <98025960+zhangpan2001@users.noreply.github.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[CPU] Dequantize feeding a strided conv is not fused, whole-tensor temp lands on the stack

1 participant