Skip to content

[Fix][Relax] Scope planned storage to if branches - #20514

Open
javierdejesusda wants to merge 1 commit into
apache:mainfrom
javierdejesusda:fix/20191-plan-memory-if-scope
Open

javierdejesusda wants to merge 1 commit into
apache:mainfrom
javierdejesusda:fix/20191-plan-memory-if-scope

Conversation

@javierdejesusda

Copy link
Copy Markdown
Contributor

Fixes #20191.

StaticPlanBlockMemory plans with one token pool per function, and its rewriter emits R.memory.alloc_storage only at the first use of a token. A token released in one if branch can be reused in the other branch or after the if, where that storage variable is not defined. The planned module is then ill-formed, and the VM segfaults for cond=False in the issue script.

The rewriter now restores its token-to-storage-variable map when it leaves a SeqExpr, so a token first used in a branch gets a new alloc_storage wherever it is reused outside that branch. Storage allocated before the if is still shared with the branches.

Tests:

  • python -m pytest -q tests/python/relax/test_transform_static_plan_block_memory.py (32 passed). Two of the three new tests fail without the change; the third pins that storage allocated before the if stays shared, and passes either way.
  • The issue script without exec_mode="bytecode" (relax.build no longer takes it), built with the default pipeline for llvm, returns x for cond=True and cond=False; before the change cond=False segfaults.
  • pre-commit run --files src/relax/transform/static_plan_block_memory.cc tests/python/relax/test_transform_static_plan_block_memory.py

Fixes apache#20191.

`StaticPlanBlockMemory` plans with one token pool per function, and its rewriter emits `R.memory.alloc_storage` only at the first use of a token. A token released in one `if` branch can be reused in the other branch or after the `if`, where that storage variable is not defined. The planned module is then ill-formed, and the VM segfaults for `cond=False` in the issue script.

The rewriter now restores its token-to-storage-variable map when it leaves a `SeqExpr`, so a token first used in a branch gets a new `alloc_storage` wherever it is reused outside that branch. Storage allocated before the `if` is still shared with the branches.

Tests:

- `python -m pytest -q tests/python/relax/test_transform_static_plan_block_memory.py` (`32 passed`). Two of the three new tests fail without the change; the third pins that storage allocated before the `if` stays shared, and passes either way.
- The issue script without `exec_mode="bytecode"` (`relax.build` no longer takes it), built with the default pipeline for llvm, returns `x` for `cond=True` and `cond=False`; before the change `cond=False` segfaults.
- `pre-commit run --files src/relax/transform/static_plan_block_memory.cc tests/python/relax/test_transform_static_plan_block_memory.py`

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.

[Bug] [Bug][Relax] StaticPlanBlockMemory can emit non-dominating storage across if branches for reshape chains

1 participant