Skip to content

[mpmd] Extract WrapBlockArgReturnsPass from UniquifyFunctionInputsOutputsPass.#1302

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_910564666
Open

[mpmd] Extract WrapBlockArgReturnsPass from UniquifyFunctionInputsOutputsPass.#1302
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_910564666

Conversation

@copybara-service
Copy link
Copy Markdown

[mpmd] Extract WrapBlockArgReturnsPass from UniquifyFunctionInputsOutputsPass.

Separates the two orthogonal concerns in UniquifyFunctionInputsOutputsPass:

  • Block-arg wrapping: ensuring block arguments returned directly by MPMD
    functions pass through a fragment (for sharding constraint attachment).
  • Return uniquification: creating unique SSA values for duplicated returns.

Introduces WrapBlockArgReturnsPass (mpmd-wrap-block-arg-returns) which handles
the first concern. The existing pass's fragment mode is simplified to only
handle duplicated returns (size > 1 filter).

The export pipeline now runs wrap before uniquify. The transfer mode used by
the sharding propagation pipeline is completely unaffected.

…putsPass.

Separates the two orthogonal concerns in UniquifyFunctionInputsOutputsPass:
- Block-arg wrapping: ensuring block arguments returned directly by MPMD
  functions pass through a fragment (for sharding constraint attachment).
- Return uniquification: creating unique SSA values for duplicated returns.

Introduces WrapBlockArgReturnsPass (mpmd-wrap-block-arg-returns) which handles
the first concern. The existing pass's fragment mode is simplified to only
handle duplicated returns (size > 1 filter).

The export pipeline now runs wrap before uniquify. The transfer mode used by
the sharding propagation pipeline is completely unaffected.

PiperOrigin-RevId: 910564666
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.

1 participant