Move softmax matcher and ErrorCheckingTrackingListener out of iree-dialects (#24466)#24543
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. |
|
i don't think the fails are related to my changes |
AGindinson
left a comment
AGindinson
left a comment
There was a problem hiding this comment.
Should the re-implemented transform_ext bits be removed within the same PR for clarity? Other than that and overall, LGTM - feels like good progress, and thanks for working on this.
The softmax StructuredOpMatcher infrastructure and ErrorCheckingTrackingListener
were relocated into GlobalOptimization/ and Codegen/Common/ in prior commits.
This deletes the now-dead originals in iree-dialects:
- Transforms/TransformMatchers.{h,cpp}
- the LinalgTransform StructuredTransformOpsExt transform-dialect extension
(.td/.h/.cpp + TableGen) and its embedded ErrorCheckingTrackingListener
- the iree-dialects CAPI (Dialects.{h,cpp}) and the Python bindings for the
extension
No compiler code includes any iree-dialects header anymore, so the deps on
IREELinalgTransformDialect / IREEDialectsTransforms / CAPI across the compiler
were stale; they are removed and the mirrored CMakeLists regenerated via
bazel_to_cmake. The ireeRegisterTransformExtensions symbol is dropped from the
public C API exports accordingly.
Progress toward retiring the iree-dialects dependency.
Signed-off-by: Alex-Wengg <hanweng9@gmail.com>
Shorten the stabilizing-max comment to a single line per review feedback: contextual rationale about how linalg.softmax decomposes belongs in the commit history (git blame), not in persistent LIT test comments where it can drift from the implementation. Signed-off-by: Alex-Wengg <hanweng9@gmail.com>
Alex-Wengg
requested review from
ScottTodd,
bangtianliu and
benvanik
as code owners
hanhanW
dismissed
their stale review
I'll leave the review to @AGindinson, thanks!
Shorten the two verbose "Negative test: ... because ..." rationale comments to one-line labels stating what each test checks, consistent with the earlier softmax_maxnumf trim. Rationale belongs in commit history, not persistent LIT test comments where it can drift. Signed-off-by: Alex-Wengg <hanweng9@gmail.com>
|
@AGindinson i have removed transform_ext and tried to reduce comments that i felt too verbose |
AGindinson
left a comment
AGindinson
left a comment
There was a problem hiding this comment.
Thanks! I had initially thought that only the corresponding functionality would be dropped, but it actually makes sense to just drop the retired code right away. I will reflect this in the squashed commit message when merging.
|
i don't think PkgCI is related to my issue linux_x64_bazel is taking quite a long time |
4 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.