Move softmax matcher and ErrorCheckingTrackingListener out of iree-dialects (#24466)#24543
Move softmax matcher and ErrorCheckingTrackingListener out of iree-dialects (#24466)#24543Alex-Wengg wants to merge 5 commits into
Conversation
Alex-Wengg
force-pushed
the
retire-iree-dialects-task1-4
branch
from
0967fa6 to
9d71a85
Compare
Alex-Wengg
requested review from
Groverkss,
IanWood1,
MaheshRavishankar,
Max191,
hanhanW,
krzysz00,
kuhar,
nirvedhmeshram and
qedawkins
as code owners
Alex-Wengg
force-pushed
the
retire-iree-dialects-task1-4
branch
2 times, most recently
from
11a1487 to
0eda000
Compare
Alex-Wengg
force-pushed
the
retire-iree-dialects-task1-4
branch
from
0eda000 to
3b008f9
Compare
|
Hi there @Alex-Wengg , small question about the From the mentioned issue, I understood that the objective was moving only the softmax matcher needed by Is keeping the full matcher helper library intentional? If not, would it make sense to trim this to |
|
@L-roro for extra context, i am still a novice when it comes to IREE so i don't have all the context but based on my additional scan, it does seem like what you have suggested is indeed correct. however i wonder if that might be out of scope ffor this PR and it could get quite messy. |
hanhanW
left a comment
hanhanW
left a comment
There was a problem hiding this comment.
This is not just a moving code task. You need to revisit the semantics and maybe rework on the matchers. That's the main reason why I filed the issue instead of doing it myself.
Alex-Wengg
force-pushed
the
retire-iree-dialects-task1-4
branch
from
74bec5d to
eb49ca5
Compare
Relocate the StructuredOpMatcher infrastructure and makeSoftmaxMatcher from llvm-external-projects/iree-dialects into the GlobalOptimization directory, which holds the only production user (RaiseSpecialOps). The code is moved verbatim (NFC); RaiseSpecialOps is repointed to the local header and the build deps are switched to the MLIR libraries the matcher needs. The now-unused matchers (reduction/matmul/convolution/pad) come along verbatim and can be trimmed in a follow-up once a build is available to validate. Progress toward retiring the iree-dialects dependency. Signed-off-by: Han || Alex <36247722+Alex-Wengg@users.noreply.github.com>
…ee-org#24466) Relocate ErrorCheckingTrackingListener into Codegen/Common (namespace mlir::iree_compiler) and repoint its users in the Common and LLVMGPU transform extensions. Drop the StructuredTransformOpsExtension registration and includes from CommonDialectRegistration and Interfaces: its transform ops (transform.iree.register_match_callbacks/match_callback/take_first/emit_remark) have no production users and their tests were removed in iree-org#24500. Build deps are updated accordingly. Progress toward retiring the iree-dialects dependency. Signed-off-by: Han || Alex <36247722+Alex-Wengg@users.noreply.github.com>
Pure reordering of include blocks in four files to satisfy the pre-commit clang-format hook (CI lint job was failing on this PR): - Codegen/Common/ErrorCheckingTrackingListener.cpp - Codegen/Interfaces/Interfaces.cpp - GlobalOptimization/TransformMatchers.cpp - GlobalOptimization/TransformMatchers.h clang-format wants llvm/* headers ordered before mlir/* within each include block (alphabetical). No removals, no behavior change. Signed-off-by: Alex-Wengg <hanweng9@gmail.com>
Replace the wholesale relocation of the iree-dialects TransformMatchers
DSL with a self-contained, native matcher in RaiseSpecialOps.cpp, per the
review feedback that this should port only what RaiseSpecialOps needs
rather than carry the generic StructuredOpMatcher framework into
GlobalOptimization.
- Add a local matchSoftmax() plus small helpers that walk the softmax
linalg-op graph directly (reduce_max -> sub -> exp -> reduce_add ->
mul/reciprocal or div), handling both implicit (projected map) and
explicit (pass-through generic) broadcasts. This is behaviorally
faithful to makeSoftmaxMatcher, including the same-source invariant.
- Delete GlobalOptimization/TransformMatchers.{h,cpp} (~3000 lines) and
drop the transform-dialect build deps that only existed for them. The
iree-dialects deps stay removed.
- Add negative lit cases (wrong max init, mismatched source) alongside
the existing softmax raising tests.
Signed-off-by: Han || Alex <36247722+Alex-Wengg@users.noreply.github.com>
The stabilizing max of a softmax can be spelled with either arith.maximumf (NaN-propagating, as emitted by e.g. StableHLO frontends) or arith.maxnumf (NaN-ignoring). The latter is what linalg.softmax itself decomposes to (SoftmaxOp::decomposeOperation and the iree-codegen-decompose-softmax pass both use arith.maxnumf), so matching only arith.maximumf made the matcher narrower than the op it raises to and missed that form. Accept both ops for the max reduction. This is strictly safer: maxnumf is the form linalg.softmax decomposes to, so raising it introduces no NaN behavior change. Add a positive lit test for the maxnumf spelling. Signed-off-by: Han || Alex <36247722+Alex-Wengg@users.noreply.github.com>
Alex-Wengg
force-pushed
the
retire-iree-dialects-task1-4
branch
from
eb49ca5 to
18d5f96
Compare
|
@hanhanW @L-roro apologies, I misunderstood the task at hand I deleted TransformMatchers.{h,cpp} and reimplemented softmax detection in RaiseSpecialOps.cpp to the best of my abilities. Dropped the iree-dialects deps as well. While reimplementing I noticed the matcher only accepted arith.maximumf, but linalg.softmax decomposes to arith.maxnumf from my understanding that would suggest the form wasn't being raised at all. I widened the matcher to accept both. I kept the existing innermost-dim-only / max-stabilized assumptions as-is not sure if those should be relaxed, let me know if that's what you had in mind. For testing I rebased onto latest main and rebuilt from scratch (incl. the recent LLVM bumps). It built cleanly, the GlobalOpt lit suite is green, and MNIST + a small StableHLO attention block both raise to linalg.softmax and run correctly end-to-end on CPU. pls lmk if this is closer to what is in mind, i will reiterate if not. |
3 participants
Progress toward #24466 (retire the
llvm-external-projects/iree-dialectsdependency). Relocates the last two pieces of code that still depended on the NFC directory. Finishing Tasks 1 & 4.
StructuredOpMatcherinfrastructure +makeSoftmaxMatcherintoGlobalOptimization/.Codegen/Common/and de-registerStructuredTransformOpsExtensionfromCommonDialectRegistrationandInterfaces. continuation of [Codegen] Remove deprecated transform.iree.match_callback tests #24500.Once this is merged, we can start removing the llvm-external-projects/iree-dialects as a whole
Testing
Built and lit tested locally on Windows 11 + MSVC 14.44 (VS 2022 Community), Ninja, CMake 3.29, Release with assertions,
IREE_TARGET_BACKEND_LLVM_CPU=ON. All changed files compile cleanly. The softmax lit testcompiler/src/iree/compiler/GlobalOptimization/test/raise_special_ops.mlirpassed. all three variants (@softmax,@softmax_no_rcp,@softmax_broadcast) match correctly through the relocatedmakeSoftmaxMatcher.