Fix RefineCallOpPattern in StablehloRefineShapes to preserve side-effecting callees - #2934
Fix RefineCallOpPattern in StablehloRefineShapes to preserve side-effecting callees#2934christopherbate wants to merge 2 commits into
Conversation
…ecting callees The constant-function replacement in RefineCallOpPattern did not check whether the callee contains side-effecting operations. This caused calls to functions with side effects (e.g. custom_call with has_side_effect=true) to be silently erased. Guard the replacement with a walk that checks for MemoryEffectOpInterface declarations.
| for (auto constAttr : constantAttrs.value()) { | ||
| constants.push_back( | ||
| ConstantOp::create(rewriter, op.getLoc(), constAttr)); | ||
| auto sideEffectResult = callee.walk([](Operation* nestedOp) { |
There was a problem hiding this comment.
nit - bool hasSideEffects = callee.walk(...).wasInterrupted(); to make it clear that interruption indicates side effect in the same line
There was a problem hiding this comment.
Also I'm not sure if this is a complete solution. If there's a call op with a nested call op I don't think this catches its side effects. We had this problem in simplification as well and made a similar method:
stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
Lines 358 to 375 in a48c8e6
A complete solution would likely require an analysis pass that maps call->side effects and does a full pass to calculate transitive side effects
There was a problem hiding this comment.
Ok let me see if there's a way I can address this by computing it up front... I already didn't like the walk because the result should be cached
## Summary - compute and cache transitive memory-effect summaries before shape refinement - erase fully constant calls only when their direct and transitive callees are memory-effect-free - retain calls that may have memory effects while propagating each constant integer-tensor result to its users - keep token and nonconstant results connected to the retained call - preserve call attributes when specializing global constant shape arguments - retain existing constant folding through memory-effect-free call chains This supersedes #2934 by covering nested func.call chains and caching function summaries. ## Testing - focused stablehlo_refine_shapes.mlir lit test (1/1 passed) - full StableHLO lit suite (272/272 passed) - ./build_tools/github_actions/lint_clang_format.sh -b upstream/main - ./build_tools/github_actions/lint_whitespace_checks.sh -b upstream/main - ./build_tools/github_actions/lint_buildifier.sh
The constant-function replacement in RefineCallOpPattern did not check
whether the callee contains side-effecting operations. This caused calls
to functions with side effects (e.g. custom_call with
has_side_effect=true) to be silently erased. Guard the replacement with
a walk that checks for MemoryEffectOpInterface declarations.