Introduce PlaceEnv for device context management

rfcs/APIs/20260214_design_for_PlaceEnv.md

@@ -0,0 +1,221 @@
+
+# 为Paddle新增支持上下文管理器与函数装饰器的设备环境切换API
+
+|任务名称 | 新增支持上下文管理器与函数装饰器的设备环境切换API |
+|---|---|
+|**提交作者** | PlumBlossom |
+|**提交时间** | 2026-02-14 |
+|**版本号** | V1.0 |
+|**依赖飞桨版本** | develop |
+|**文件名** | 20260214_design_for_PlaceEnv.md |
+
+# 一、概述
+## 1、相关背景
+在深度学习模型的开发与迁移过程中，开发者常常需要精确控制代码在不同硬件设备（如CPU、GPU、NPU等）上的执行部分。例如，在将ECDFormer等模型从PyTorch迁移至Paddle时，数据集的加载和预处理阶段包含了大量`paddle.to_tensor`、`paddle.ones`、`paddle.reshape`等操作。这些操作若在GPU上执行，会造成显存的浪费，且对于某些操作而言，CPU的执行效率更高。虽然Paddle提供了`paddle.set_device`来全局设置设备，但在复杂的数据流水线中，开发者可能希望临时将某段代码或某个函数强制在指定设备（如CPU）上运行，并在执行完毕后恢复原设备上下文。当前Paddle的`paddle.device.device`上下文管理器仅支持with语句，且功能较为单一，无法直接作为装饰器作用于函数，给开发者带来不便。
+
+## 2、功能目标
+实现一个新的API,命名为`paddle.PlaceEnv`，使其具备以下能力：
+1.  **上下文管理器**：支持`with PlaceEnv(device):`语法，在该代码块内部，所有Paddle算子的默认设备切换为指定的设备（如CPUPlace、CUDAPlace、CustomPlace等），代码块执行完毕后，自动恢复至进入前的设备设置。
+2.  **函数装饰器**：支持`@PlaceEnv(device)`语法，当装饰一个函数时，该函数体内的所有Paddle算子均在指定的设备上下文中执行，函数执行完毕后退回原设备。
+3.  **设备类型扩展**：不仅支持Paddle内置的CPU和GPU，还必须支持通过`Place`接口传入的自定义硬件，例如`paddle.CustomPlace('my_hardware:0')`，以确保对国产硬件（如寒武纪、昇腾等）的兼容性。
+
+## 3、意义
+- **提升代码健壮性与效率**：开发者可以轻松地将数据预处理等CPU友好型操作固定在CPU上执行，避免因`set_device`位置不当导致意外的GPU显存占用和性能损耗。
+- **增强代码复用性**：通过装饰器，可以方便地将一个已经写好的、未显式指定设备的函数（例如某个数据转换函数）快速适配到目标设备上，无需修改函数内部代码。
+- **对齐并超越业界方案**：当前PyTorch的`torch.cuda.device`和Paddle的`paddle.device.device`均只提供上下文管理器功能。本设计增加装饰器功能，是对现有生态的补充和优化，提供了更灵活的编程范式。
+- **支持硬件多元化**：设计之初即考虑自定义硬件Place，确保API能够平滑支持Paddle生态内的所有硬件后端，符合当前AI芯片多样化的发展趋势。
+
+# 二、飞桨现状
+目前，Paddle框架中与设备切换相关的接口主要有：
+1.  `paddle.set_device()`：全局设置默认设备。其影响是全局的，无法细粒度地控制局部代码的设备环境。
+2.  `paddle.device.device()`：这是一个上下文管理器，允许在`with`语句块中临时切换设备。其[官方文档](https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/cuda/device_cn.html)表明它仅支持作为上下文管理器使用，不支持作为函数装饰器。
+
+因此，当前无法通过装饰器便捷地为一个完整的函数指定设备上下文。虽然有替代方案（如在函数内部手动调用`set_device`并在函数结束前恢复，或手动将函数内容包裹在`with`块中），但这些方式侵入性强、代码冗余且易出错。
+
+# 三、业内方案调研
+我们调研了主流框架及相关库的实现：
+1.  **PyTorch (`torch.cuda.device`)**：提供了上下文管理器，用于临时选择CUDA设备。其实现基于Python的上下文管理器协议（`__enter__`和`__exit__`），通过修改线程局部状态来改变当前设备。**不支持函数装饰器**。
+2.  **TensorFlow (`tf.device`)**：提供了强大的设备上下文管理器，可以接受设备名称字符串。它也支持作为装饰器使用，例如：
+    ```python
+    @tf.device('/GPU:0')
+    def matmul_on_gpu():
+        # 此函数内的操作将在GPU:0上执行
+        ...
+    ```
+    TensorFlow的实现依赖于其图构建阶段的设备放置机制，通过修改线程的上下文堆栈来实现。
+3.  **JAX (`jax.devices` 和 `jax.default_device`)**：JAX提供了`jax.default_device`上下文管理器，用于临时设置默认执行设备。它基于Python的`contextlib.contextmanager`实现，但主要影响JIT编译后的内核放置。
+
+**未来趋势**：随着硬件异构计算的普及，框架层对设备管理的灵活性要求越来越高。提供更细粒度、更便捷的设备控制手段是框架发展的趋势。装饰器作为一种声明式的编程模式，因其非侵入性而受到开发者青睐。
+
+# 四、对比分析
+| 方案 | 上下文管理器支持 | 函数装饰器支持 | 自定义硬件支持 | 实现复杂度 | 优势 | 劣势 |
+| :--- | :--- | :--- | :--- | :--- | :--- | :--- |
+| **PyTorch方案** | ✅ | ❌ | 有限（主要CUDA） | 中 | 稳定可靠，用户熟悉 | 功能单一，无法覆盖自定义硬件和装饰器场景 |
+| **TensorFlow方案** | ✅ | ✅ | 支持（设备字符串） | 高（依赖图机制） | 功能强大，灵活度高 | 实现机制与Paddle动态图模式差异较大，直接借鉴难度高 |
+| **Paddle现状** | ✅ | ❌ | 支持（通过Place） | 低 | 已具备良好的Place抽象 | 缺少装饰器支持，使用不够灵活 |
+| **本设计方案** | ✅ | ✅ | **支持** | 中 | **在Paddle现有基础上扩展，兼容性好，功能完备** | 需要新增API，需考虑与现有`paddle.device.device`的兼容和演进关系 |
+
+**结论**：在Paddle现有良好的`Place`抽象和上下文管理器基础上，借鉴TensorFlow的装饰器思想，实现一个功能更完备的设备环境切换API，是收益最高、对用户最友好的方案。
+
+# 五、设计思路与实现方案
+
+## 1、主体设计思路与折衷
+### 整体全貌（初步方案）
+本设计位于Paddle框架基础架构层，核心是修改设备管理模块的上下文栈。当用户通过`with PlaceEnv(place):`或`@PlaceEnv(place)`进入一个设备环境时，一个新的设备上下文被推入线程局部栈中。后续Paddle所有算子创建Tensor或执行时，会从该栈顶获取当前应使用的默认设备Place。退出时，该上下文被弹出，恢复至之前的设备环境。
+
+<img src=https://github.com/user-attachments/assets/6d8c4e67-8c03-4851-ba38-c4c9e267399e></img>
+
+*(图：PlaceEnv执行流程)*
+
+### 主体设计具体描述
+我们将创建一个新的类`paddle.PlaceEnv`。该类同时实现上下文管理器协议（`__enter__`, `__exit__`）和装饰器协议（`__call__`）。
+- **修改位置**：待定，**等待评审人员建议 🤔**。
+- **核心机制**：利用Python的`threading.local()`维护一个线程局部的设备栈。栈顶元素即为当前线程的默认设备Place。
+
+### 主体设计选型考量
+- **为什么新增类而非修改原有`device`类？** 原有`paddle.device.device`功能明确，直接修改可能引入兼容性风险，影响现有用户。新增`PlaceEnv`可以提供更丰富的功能（如装饰器），同时保持与旧API的共存，并逐步引导用户迁移。
+- **为什么选择线程局部存储？** Paddle框架内部已经广泛使用线程局部变量来管理设备上下文，这是实现设备切换的标准且高效的方式，可以保证在复杂的多线程数据加载场景下，每个线程的设备上下文是独立的。
+
+## 2、关键技术点/子模块设计与实现方案
+### 技术点1：线程局部设备栈管理
+**数据结构**：在`paddle.device`模块内部，定义一个`_thread_local = threading.local()`，并为其设置一个初始属性`_thread_local.device_stack = []`。
+**入栈操作**：
+```python
+def push_device(place):
+    # 获取当前默认place（栈顶或全局设置）
+    current_place = get_current_device()
+    # 将当前place和新place的信息（用于恢复）与新place一起推入栈
+    # 栈元素可以是一个元组 (previous_place, new_place)
+    stack = get_device_stack()
+    stack.append( (current_place, place) )
+    # 设置新place为默认
+    _set_internal_device(place)
+```
+**出栈操作**：
+```python
+def pop_device():
+    stack = get_device_stack()
+    if stack:
+        previous_place, _ = stack.pop()
+        _set_internal_device(previous_place)
+```
+
+### 技术点2：`PlaceEnv`类的实现
+```python
+class PlaceEnv:
+    def __init__(self, place):
+        # 将输入的设备描述（如'cpu', 'gpu:0', CustomPlace对象）统一转化为paddle.Place对象
+        self.place = _convert_to_place(place)
+
+    def __enter__(self):
+        # 将self.place推入设备栈，并设置为当前设备
+        push_device(self.place)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # 从设备栈弹出，恢复至之前的设备
+        pop_device()
+
+    def __call__(self, func):
+        # 实现装饰器逻辑
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with self:
+                return func(*args, **kwargs)
+        return wrapper
+```
+
+### 技术点3：获取当前设备函数`get_current_device()`
+需要提供一个内部函数来获取当前有效的默认设备。其逻辑为：若线程局部栈不为空，则返回栈顶元素的设备；否则，返回`paddle.get_device()`或全局默认Place。
+
+## 3、主要影响的模块接口变化
+### 直接接口变化
+- **新增接口**：`paddle.device.PlaceEnv(place)`。place参数支持`int`, `str`, `paddle.Place`。
+  - 使用示例1（上下文管理器）：
+    ```python
+    import paddle
+    print(paddle.device.get_device()) # 假设输出 'gpu:0'
+    with paddle.device.PlaceEnv('cpu'):
+        # 此块内所有操作在CPU上执行
+        tensor = paddle.ones([10, 10]) # 实际在CPU上创建
+        print(tensor.place) # 输出 CPUPlace
+    print(paddle.device.get_device()) # 恢复为 'gpu:0'
+    ```
+  - 使用示例2（函数装饰器）：
+    ```python
+    import paddle
+
+    @paddle.device.PlaceEnv(paddle.CPUPlace())
+    def load_and_preprocess_data():
+        # 整个函数在CPU上执行
+        data = paddle.to_tensor([1,2,3])
+        return data
+
+    processed_data = load_and_preprocess_data() # 函数执行在CPU，但外部设备不变
+    print(processed_data.place) # 输出 CPUPlace
+    ```
+  - 使用示例3（自定义硬件）：
+    ```python
+    import paddle
+    # 假设已经注册了名为'my_hardware'的自驱动硬件
+    with paddle.device.PlaceEnv('my_hardware:0'):
+        tensor = paddle.ones([10, 10]) # 在自定义硬件上执行
+        print(tensor.place) # 输出 CustomPlace(my_hardware, 0)
+    ```
+
+### 对框架各环节的影响排查
+- **网络定义**：无直接影响。装饰器或with块包裹整个网络定义时，网络将在指定设备上构建。
+- **底层数据结构**：核心影响点是`Tensor`的`place`属性。本设计确保在上下文内创建的Tensor的place与上下文指定的一致。
+- **OP**：OP执行时会从上下文中获取设备信息，本设计保证了`get_current_device`等内部函数返回值的正确性。
+- **数据IO**：无直接影响，但鼓励用户在数据加载部分使用此API将数据强制放在CPU，是主要应用场景。
+- **执行**：在执行引擎层面，设备上下文栈的管理是透明的，不影响现有执行流程。
+- **分布式**：在多卡或多节点训练中，每个进程/线程的设备栈是独立的，因此不会相互干扰。
+- **模型保存**：保存模型时，`state_dict`中的Tensor会保留其原有的place信息，但加载时可以自由放置，故无影响。
+- **预测部署**：预测库中如果使用了Paddle的Python API，同样可以受益于此功能。
+
+# 六、测试和验收的考量
+- **自测方案**：
+    1.  **单元测试**：覆盖所有设备类型（CPU/GPU/自定义Place），测试嵌套的with语句、with和装饰器混合使用、多线程场景下设备上下文的正确隔离与恢复。
+    2.  **功能测试**：编写一个包含`paddle.to_tensor`, `paddle.ones`等操作的函数，分别使用装饰器和不使用装饰器，验证其返回Tensor的place属性是否正确。
+- **CE（兼容性测试）**：
+    1.  确保新增API与现有`paddle.device.device`、`paddle.set_device`同时使用时行为符合预期，且互不破坏。
+    2.  验证在无GPU、有GPU、有自定义硬件的多种环境下，API都能正确工作或给出清晰的错误提示。
+- **验收度量**：
+    1.  所有测试用例通过率100%。
+    2.  文档示例清晰，能够指导用户完成从上下文管理器到装饰器的平滑迁移。
+    3.  社区反馈积极，无因本API引入的bug报告。
+
+# 七、影响面
+## 对用户的影响
+用户获得了一个更强大的工具来管理代码的设备环境，编写更清晰、更高效的代码。特别是对于从PyTorch迁移模型的开发者，可以轻松解决数据预处理阶段的设备控制问题。这是一个纯粹的增强功能，对不使用此API的现有代码无任何影响。
+
+## 对二次开发用户的影响
+- **新增API**：`paddle.PlaceEnv`会暴露给二次开发用户，作为他们进行设备管理的高级工具。
+- **内部机制**：`_thread_local.device_stack`相关的内部接口不建议公开，但开发者若需深度定制设备管理逻辑，可以此为参考。
+
+## 对框架架构的影响
+在现有设备管理模块上增加了一层薄薄的封装，对核心架构无侵入式修改。代码结构清晰，易于维护。
+
+## 对性能的影响
+设备上下文的入栈和出栈操作是常数时间复杂度的Python级操作，对整体性能影响可以忽略不计。算子执行时获取当前设备的开销与原有机制一致，无额外性能损耗。
+
+## 对比业内深度学习框架的差距与优势的影响
+- **缩小差距**：补齐了Paddle在设备管理灵活性上与TensorFlow的部分差距（装饰器支持）。
+- **建立优势**：通过统一的`Place`抽象，Paddle在支持多种异构硬件方面的体验优于主要对标框架（PyTorch的`torch.cuda.device`与硬件绑定较深）。本设计充分发扬了这一优势。
+
+## 其他风险
+- **命名冲突**：`PlaceEnv`可能与其他第三方库重名，但概率较低，且在`paddle`命名空间下是安全的。
+- **滥用风险**：建议在文档中强调此API应用于局部、临时的设备切换，不建议用于全局设备设置，以免造成代码混乱。
+
+# 八、排期规划
+- 待定
+
+# 名词解释
+- **Place**：Paddle中表示设备位置的抽象，如`CPUPlace()`, `CUDAPlace(0)`, `CustomPlace()`。
+- **上下文管理器**：实现了`__enter__`和`__exit__`方法的对象，用于`with`语句中，管理资源的进入和退出。
+- **装饰器**：一种Python语法，允许在不修改函数体的情况下，给函数增加额外的功能。
+
+# 附件及参考资料
+1.  Paddle官方文档 - [paddle.device.device](https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/cuda/device_cn.html)
+2.  Python官方文档 - [contextlib](https://docs.python.org/3/library/contextlib.html) - 用于创建上下文管理器和装饰器的工具。
+3.  Python官方文档 - [functools.wraps](https://docs.python.org/3/library/functools.html#functools.wraps) - 用于定义装饰器时保留原函数元信息。
+