iree-org / iree-org/iree

Move DecomposeSoftmax to GlobalOptimization.

Open
#17,469 28 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
3.9k
Forks
1k
Avg merge
4d 16h
Merged PRs (30d)
47

Description

Currently, DecomposeSoftmax is under Codegen/. Here's a link to it, specifically to a helpful comment explaining what it rewrites softmax to:
https://github.com/iree-org/iree/blob/1316c92341ec5aa320f44ad4c8be9771b20e61e7/compiler/src/iree/compiler/Codegen/Common/DecomposeSoftmax.cpp#L19-L36

Problems with that:
* This requires an intermediate buffer on the stack at least between steps 2 and 3 of the above-linked comment (already in the 1D case) and potentially at every step (N-D general case).
* This prevents softmax from benefiting from generic fusion logic: since it is still a softmax named op at the time of dispatch region formation, only a softmax-specific piece of logic can fuse it. That's unnecessary complexity, since softmax will later be decomposed to generic ops anyway.

The main action to be taken here is to move DecomposeSoftmax to the appropriate point in Flow, before dispatch region formation.

Then, tweak the dispatch fusion logic to achieve the following goal: no buffers larger than some small fixed size should be passed on the stack, so whenever a potentially large buffer is passed between steps in the above-linked comment, that should be a dispatch region boundary.

In the 1-dimensional case (which is the most common case), that means that steps 1. and 2. can go in one dispatch, and then steps 3. and 4. can go in another. The buffer produced by step 2 can't go on the stack.

In the higher-dimensional case, potentially every step might have to be its own dispatch, but I suppose that a common sub-case of that might be when some inner dimensions are static and small, in which case it boils down to the 1-dimensional case.

Finally, perhaps as a follow-up unblocked by this groundwork, ensure that suitable producers are fused into the first dispatch (ie the one typically containing steps 1 and 2 from the description).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.