[Feature]: 静态图支持配置Cell输出做重计算

### 🚀 背景描述

重计算是平衡内存占用与计算性能的核心技术。通过在反向传播时重新执行前向算子，可以避免保存大尺寸的中间激活值。

在之前的实现中，静态图重计算功能缺乏对 Cell 输出是否参与重计算的精细控制机制，某些大模型训练场景无法更进一步地减少显存占用。动态图此前已支持配置Cell的输出做重计算，故本次需要在静态图的重计算功能中支持。

### 设计思路

## 1 总体架构
重计算标志的生命周期从 Python 侧 API 开始，经过前端编译，最终影响前端 IR 的优化行为：

```
┌───────────────────────────┐       ┌───────────────────────────┐      ┌───────────────────────────┐
│   Python 层 (Cell API)    │       │   前端 IR 解析 (Parse)     │      │   前端 IR 优化 (Pass)      │
├───────────────────────────┤       ├───────────────────────────┤      ├───────────────────────────┤
│ ms.recompute(cell,        │       │ composite.cc:             │      │ recompute.cc:             │
│   output_recompute=True)  │ ───▶ │   fg->set_attr(           │ ───▶│   IsRecomputeCell()       │
└───────────────────────────┘       │     "output_recompute")   │      │   AddNewPrimalNode()      │
                                    └───────────────────────────┘      └───────────────────────────┘
```
## 2 核心变更

### 2.1 基础常量与存储

**文件**: `mindspore/core/include/ir/func_graph_flag.h`

- **修改前**: `const char FUNC_GRAPH_OUTPUT_NO_RECOMPUTE[] = "output_no_recompute";`
- **修改后**: `const char FUNC_GRAPH_FLAG_OUTPUT_RECOMPUTE[] = "output_recompute";`

**存储逻辑**:
统一使用 `has_attr` 和 `get_attr` 进行访问。如果 `attr` 存在且值为 `true`，则表示该子图的输出允许被后续算子重计算。

### 2.2 Python 侧重构

**文件**: `mindspore/python/mindspore/nn/cell.py`

1. **`_recompute` 方法简化**：
   - 移除 `mode` 参数，统一使用 `output_recompute` 参数
   - 简化 `_set_recompute_scope` 方法，移除 `mode` 相关逻辑

2. **标志传递**：
   - 将 `output_recompute` 布尔值直接通过 `add_flags` 传递给 C++
   - 子 Cell 默认继承重计算属性

```python
def _recompute(self, output_recompute=False):
    """
    Set the cell recomputed.
    """
    ...
    self._set_recompute_scope()
    self.add_flags(output_recompute=output_recompute)
    for cell in self.cells():
        cell._recompute(True)  # 子 Cell 默认继承重计算属性
```

### 2.3 C++ 核心逻辑增强
Cell 输出做重计算，只要 Cell 输出的属性不强制打 `recompute: False` 的attr，Inline 后重计算的机制天然能对接上，因此主要适配的是 Inline 前重计算的机制。

**文件**: `mindspore/ccsrc/frontend/optimizer/irpass/recompute.cc`

引入了 `recompute_cell_output` 逻辑，用于支持 **合并相邻重计算**。

1. **`IsRecomputeCell`**: 扩展为支持返回输出重计算状态
2. **`AddNewPrimalNode` 的条件分支**:
   - 如果 `output_recompute=false`（默认）：不管使用者是不是重计算节点，输入都会连接到新的 Cell `Primal` 节点，确保 Cell 的输出被保存
   - 如果 `output_recompute=true`：当使用者也是一个重计算节点时，输入不连接到新的 Cell `Primal` 节点，从而允许两个相邻的重计算 Cell 共享激活值，进一步节省内存

```cpp
// 核心判断逻辑
bool change = AddNewPrimalNode(manager, fg, node, new_primal,
                               recompute_cell && !recompute_cell_output,
                               &origin_to_new_primal);
```

### 2.4 IR Pass 适配

**文件**: `mindspore/ccsrc/frontend/optimizer/irpass/recompute_prepare.h`

在 Inline 之前的准备阶段：
- 从 `has_flag/erase_flag` 切换为 `has_attr/get_attr/erase_attr`
- 根据 `output_recompute` 的值决定是否在输出节点上添加 `kAttrRecompute=false` 属性

## 3. 实现细节

### 3.1 重计算标志传递流程

```
┌──────────────────────────────────────────────────────────┐
│          用户调用 Cell.recompute(output_recompute=True)   │
└────────────────────────────┬─────────────────────────────┘
                             │
                             ▼
┌──────────────────────────────────────────────────────────┐
│          内部方法 Cell._recompute(output_recompute)       │
├──────────────────────────────────────────────────────────┤
│ 1. _set_recompute_scope(): 设置 'recompute_' scope前缀    │
│ 2. Cell.add_flags(output_recompute=output_recompute)     │
└────────────────────────────┬─────────────────────────────┘
                             │
                             ▼
┌──────────────────────────────────────────────────────────┐
│          前端解析器 (Parser/Composite)                    │
├──────────────────────────────────────────────────────────┤
│ 1. 识别 Python 标志并转换为 C++ 属性                       │
│ 2. 调用 FuncGraph::set_attr("output_recompute", value)   │
└────────────────────────────┬─────────────────────────────┘
                             │
                             ▼
┌──────────────────────────────────────────────────────────┐
│          前端 Recomputation Pass                         │
├──────────────────────────────────────────────────────────┤
│ 1. 遍历图节点时通过 get_attr 读取标志值                    │
│ 2. 根据标志决策节点生成逻辑                                │
└──────────────────────────────────────────────────────────┘
```

### 3.2 C++ Pass 核心决策流程

```
           开始遍历 K 图调用节点 (k_fg_caller)
                       │
                       ▼
           获取 "output_recompute" 属性值
                       │
            ┌──────────┴──────────┐
            ▼                     ▼
    [output_recompute == false]  [output_recompute == true]
    (默认: 输出不重计算)          (开启输出重计算)
            │                     │
            │            user节点是否也是重计算节点?
            │            ┌────────┴────────┐
            │            ▼ (Yes)           ▼ (No)
            │       共享激活值              │
            │       跳过创建新 Primal       │
            │       实现合并重计算          │
            │            │                 │
            ▼            ▼                 ▼
    ┌──────────────────────────────────────────────────┐
    │        AddNewPrimalNode(): 为前向user生成连接      │
    └────────────────────────┬─────────────────────────┘
                             │
                             ▼
    ┌──────────────────────────────────────────────────┐
    │        AddDependNodes(): 插入依赖边确保时机        │
    └────────────────────────┬─────────────────────────┘
                             │
                             ▼
                    完成该节点的重计算处理
```

### 3.3 关键代码片段

#### 3.3.1 IsRecomputeCell 判断逻辑

```cpp
bool IsRecomputeCell(const FuncGraphPtr &k_fg, bool *recompute_cell_output = nullptr) {
  if (k_fg == nullptr) {
    return false;
  }
  auto primal_iter = k_fg->transforms().find("primal");
  if (primal_iter == k_fg->transforms().end()) {
    MS_LOG_WITH_NODE(EXCEPTION, k_fg->return_node())
      << "The k_fg: " << k_fg << " should have primal part.";
  }
  auto primal_fg = primal_iter->second.func_graph();
  if (primal_fg == nullptr) {
    return false;
  }
  if (!primal_fg->has_attr(FUNC_GRAPH_FLAG_OUTPUT_RECOMPUTE)) {
    MS_LOG(ERROR) << "The primal graph " << primal_fg->ToString()
                  << " should have output recompute attribute.";
    return false;
  }
  if (recompute_cell_output != nullptr) {
    *recompute_cell_output = GetValue<bool>(primal_fg->get_attr(FUNC_GRAPH_FLAG_OUTPUT_RECOMPUTE));
  }
  return true;
}
```

#### 3.3.2 AddNewPrimalNode 核心逻辑

```cpp
bool AddNewPrimalNode(const FuncGraphManagerPtr &manager, const FuncGraphPtr &fg,
                     const AnfNodePtr &origin_primal, const AnfNodePtr &new_primal,
                     bool not_recompute_cell_output,
                     std::unordered_map<AnfNodePtr, AnfNodePtr> *origin_to_new_primal) {
  bool changed = false;
  auto node_users = manager->node_users()[origin_primal];
  for (auto &node_and_idx : node_users) {
    auto user = node_and_idx.first;
    MS_EXCEPTION_IF_NULL(user);
    // 处理 TupleGetItem、Depend、MakeTuple 等中间节点
    // ...
    // Set edge to not recomputed primal nodes.
    if (not_recompute_cell_output || (!IsRecomputeKGraphCaller(user) && !IsGradNode(user))) {
      MS_LOG(DEBUG) << "Set edge to user: " << user->DebugString()
                    << ", new primal: " << new_primal->DebugString();
      manager->SetEdge(user, node_and_idx.second, new_primal);
      changed = true;
    }
  }
  return changed;
}
```

## 4. 详细对比

### 涉及到的对外API

`mindspore.nn.Cell.recompute()`
 `mindspore.common.recompute()`

由于动态图已支持配置 Cell 输出也做重计算，已加上 `output_recompute` 参数，所以本次修改不涉及对外 API 的变更，只需要在静态图的逻辑中对接 `output_recompute` 参数即可。

### 与其他模块的相关性描述

NA

### 测试设计与测试计划

1、Jit下执行grad场景的重计算：
tests/st/compiler/recompute/test_recompute_with_jit.py
```
"""
Feature: Jit下求导且对Cell设置重计算
Description: 每个block使用Cell的recompute接口设置重计算且output_recompute参数设置为True
Expectation: 用例执行不报错且显存占用符合预期
"""
@arg_mark(plat_marks=['platform_ascend'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_recompute_cell_recompute_with_jit_with_output_recompute():
    run_testcase("test_recompute_block_recompute_with_jit_with_output_recompute", 88)
```

2、Jit下执行grad且带LazyInline场景下重计算：
tests/st/compiler/recompute/test_recompute_with_lazy_inline.py
```python
"""
Feature: Jit下求导且带Lazy Inline场景下对Cell设置重计算
Description: 每个block使用Cell的recompute接口设置重计算且output_recompute参数设置为True
Expectation: 用例执行不报错且显存占用符合预期
"""
@arg_mark(plat_marks=['platform_ascend'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_recompute_cell_recompute_with_output_recompute():
    run_testcase("test_recompute_block_recompute_with_output_recompute", 88)

"""
Feature: Jit下求导且带Lazy Inline场景下对Cell设置重计算
Description: 每个block使用ms.common.recompute接口设置重计算且output_recompute参数设置为True
Expectation: 用例执行不报错且显存占用符合预期
"""
@arg_mark(plat_marks=['platform_ascend'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_recompute_cell_recompute_func_api_with_output_recompute():
    run_testcase("test_recompute_block_recompute_func_api_with_output_recompute", 88)

"""
Feature: Jit下求导且带Lazy Inline场景下对Cell和Primitive设置重计算
Description: 每个block使用Cell和Primitive的recompute接口设置重计算且Cell的重计算接口output_recompute参数设置为True
Expectation: 用例执行不报错且显存占用符合预期
"""
@arg_mark(plat_marks=['platform_ascend'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_recompute_cell_and_op_recompute_with_output_recompute():
    run_testcase("test_recompute_cell_and_op_recompute_with_output_recompute", 61)
```

3、GradJit场景下重计算：
tests/st/compiler/recompute/test_recompute_under_gradjit.py
```
"""
Feature: GradJit场景下对Cell设置重计算
Description: block使用Cell的recompute接口设置重计算且output_recompute参数设置为True
Expectation: IR上输出算子的数量符合预期
"""
@arg_mark(plat_marks=['cpu_linux'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_recompute_block_recompute_with_output_recompute():
    class NetRecompute(Cell):
        def __init__(self):
            super().__init__()
            self.block = LeNet()
            self.block.recompute(output_recompute=True)

@jit
        def construct(self, x):
            out = self.block(x)
            out = ops.Abs()(out)
            return out

save_graphs_path = "./test_recompute_block_recompute_with_output_recompute"
    context.set_context(save_graphs=True, save_graphs_path=save_graphs_path)

x = Tensor(np.ones([32, 1, 32, 32]).astype(np.float32) * 0.01)
    net = NetRecompute()
    grad_net = Grad(net)
    grad_net(x)

para = '= BiasAdd(%'
    output = subprocess.check_output(
        ["grep -r '%s' %s | wc -l" % (para, os.path.join(save_graphs_path, "opt_backward_[0-9]*.ir"))],
        shell=True)
    out = str(output, 'utf-8').strip()
    assert out == "3"

if os.path.exists(save_graphs_path):
        shutil.rmtree(save_graphs_path)
```

### 其他信息

## 改进收益

增强了接口灵活性，通过精细控制 Cell 输出重计算，可以合并相邻重计算 Cell，避免不必要的激活值保存，可进一步减少大模型训练在流水线并行等场景下的内存峰值。

GVP MindSpore / mindspore
关闭

内容风险标识

评论 (2)

GVPMindSpore / mindspore关闭

内容风险标识