已合并
Fix mfusion FX-to-torch-mlir conversion for torch.ops.higher_order.run_and_save_rng_state #38582
shuqian0创建于 6月15日
Fix mfusion FX-to-torch-mlir conversion for torch.ops.higher_order.run_and_save_rng_state #38582
已合并
S
shuqian0
6月15日
6月15日
# 【合入来源】
> **如有社区issue,请关联issue链接**\
> **请勿携带内部流程信息(需求链接、问题单、内部issue等)**
- [ ] 需求
- [x] 问题单
- [ ] issue/工单
- [ ] 重构优化
- [ ] 资料更新
# 【修改方案】
[[Feature]: 26.1.0版本图模式功能增强](https://gitcode.com/Ascend/pytorch/issues/1978)
修复 mfusion 中 `torch.ops.higher_order.run_and_save_rng_state` 算子从FX 到 torch-mlir 的转换报错。
失败的案例是:
```python
run_and_save_rng_state = torch.ops.higher_order.run_and_save_rng_state(...)
rng_state = run_and_save_rng_state[0]
dropout_out = run_and_save_rng_state[1][0]
```
torch-mlir 不直接实现这个 HOP(Higher Order Primitive),其导入器将多输出操作建模为扁平化的 MLIR 结果。本PR将不支持的目标通过 mfusion opaque 自定义操作路由,同时保留足够的结构以便在 MLIR 往返转换后重建原始 FX 图。
## 变更内容
- 仅将 torch-mlir 支持的 HOP 视为可直接导入。不支持的 HOP 现在使用现有的 opaque 往返转换路径,而不是进入 torch-mlir 的 `_import_hop`。
- 将 opaque tuple 结果元数据扁平化为 MLIR 兼容的多结果,并在 `opaque_registry.Payload` 中记录每个叶子结果的路径。
- 在 torch-mlir 导入之前,将 opaque tuple 结果的 getitem 用户重写为从扁平化的 opaque 结果中读取。
- 在 MLIR 到 FX 导出期间重建原始 opaque/HOP 节点和嵌套的 getitem 路径,使用前缀节点复用以避免重复的嵌套 getitems。
- 为不支持的单叶子 tuple 输出(如 `(Tensor,)`)添加明确的错误提示,这种输出无法在不改变语义的情况下由当前的扁平化方案表示。
# 【资料变更】
> 请确认是否涉及资料变更。如涉及,需要在PR中体现,并简要说明修改内容。如不涉及,需填写“不涉及”
# 【接口变更】
> 请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及,需要详细说明接口以及对应的变更内容,同时需要在资料中体现。如不涉及,需填写“不涉及”
# 【功能验证】
> 说明测试场景,测试方法。如果本次测试方式与常规单元测试不同,请详细说明您的测试步骤\
> 新增/变更内容是否已新增/适配UT测试用例看护,并补充测试自验证截图
# 【CheckList】
> PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]
- [x] 代码注释完备,正确记录错误日志
- [x] 代码实现进行了返回值、空指针等校验
- [x] PR标题正确使用类型标签,如:feat、fix、refactor、docs、test等
- [x] PR持续集成流水线(CI)执行通过,代码检查无异常


Pull Request已成功合入, 合并人@ascend-robot
(感谢 shuqian0 的贡献)ascend-robot
6月15日 评论:
6月15日 评论:
6月15日 添加了label:ascend-cla/yes
ascend-robot
6月15日 评论:
6月15日 评论:
Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| test | ✅ zichun_ye, hbhu_bin (2/2) | ✅ zichun_ye (1/1) |
| torch_npu/_inductor | ✅ zichun_ye, hbhu_bin (2/2) | ✅ zichun_ye (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
shuqian0, thanks for your pull request. All authors of the commits have signed the CLA. 👍


6月15日 添加了label:needs-issue
此处折叠了68条消息 查看更多
6月23日 添加了label:approved
6月23日 添加了label:lgtm
6月23日 合入了pull request
ascend-robot
6月23日 评论:
6月23日 评论:
流水线 pytorch_gitcode_PR_multiVersion#11110 [ commitID:1c2ed023 ] 已完成

