From 520aac039f80b12dee0bb2ee4478828a8cad50b7 Mon Sep 17 00:00:00 2001 From: megemini Date: Thu, 30 Oct 2025 14:06:13 +0800 Subject: [PATCH 1/3] =?UTF-8?q?=E3=80=90Hackathon=209th=20No.109=E3=80=91?= =?UTF-8?q?=E5=9F=BA=E4=BA=8E=20Setuptools=2080+=20=E7=89=88=E6=9C=AC?= =?UTF-8?q?=E8=87=AA=E5=AE=9A=E4=B9=89=E7=AE=97=E5=AD=90=E6=9C=BA=E5=88=B6?= =?UTF-8?q?=E9=80=82=E9=85=8D=E8=AE=BE=E8=AE=A1=E6=96=87=E6=A1=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...design_for_setuptools80_custom_operator.md | 316 ++++++++++++++++++ 1 file changed, 316 insertions(+) create mode 100644 rfcs/APIs/20251030_api_design_for_setuptools80_custom_operator.md diff --git a/rfcs/APIs/20251030_api_design_for_setuptools80_custom_operator.md b/rfcs/APIs/20251030_api_design_for_setuptools80_custom_operator.md new file mode 100644 index 000000000..fa78ead1d --- /dev/null +++ b/rfcs/APIs/20251030_api_design_for_setuptools80_custom_operator.md @@ -0,0 +1,316 @@ +# paddle.utils.cpp_extension 基于 Setuptools 80+ 版本自定义算子机制适配设计文档 + +| API名称 | paddle.utils.cpp_extension 基于 Setuptools 80+ 版本自定义算子机制适配设计文档 | +|---|---| +|提交作者 | megemini | +|提交时间 | 2025-10-30 | +|版本号 | V1.0 | +|依赖飞桨版本 | develop版本 | +|文件名 | 20251030_api_design_for_setuptools80_custom_operator.md
| + + +# 一、概述 + +## 1、相关背景 + +关联任务：https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_9th/%E3%80%90Hackathon_9th%E3%80%91%E4%B8%AA%E4%BA%BA%E6%8C%91%E6%88%98%E8%B5%9B%E2%80%94%E6%A1%86%E6%9E%B6%E5%BC%80%E5%8F%91%E4%BB%BB%E5%8A%A1%E5%90%88%E9%9B%86.md#no109-基于-setuptools-80-版本自定义算子机制适配 + +关联 PR：https://github.com/PaddlePaddle/Paddle/pull/76008 + +PaddlePaddle 目前对于自定义 C++ 算子的实现是基于 setuptools 做了一些 patch，在 bdist_egg 阶段通过 patch write_stub 实现的，然而在 setuptools 80+，被 patch 的逻辑在 install command 不会被走到（于 pypa/setuptools#2908 移除），因此我们希望基于 setuptools 80+ 对自定义 C++ 算子进行适配，确保自定义 C++ 算子在 setuptools 80+ 是可用的。 + +## 2、功能目标 + +适配 Setuptools 80+ 版本，确保自定义算子在新版本 Setuptools 下能够正常编译、安装和运行，并保持向后兼容性，确保在旧版本 Setuptools 下仍能正常工作。 + +## 3、意义 + +确保自定义算子在新版本 Setuptools 下能够正常编译、安装和运行。 + +# 二、飞桨现状 + +PaddlePaddle 目前的自定义算子机制主要通过 `paddle.utils.cpp_extension` 模块实现，依赖 `write_stub` 机制生成 Python API 文件，该机制在 Setuptools 80+ 中不再自动触发。 + +# 三、业内方案调研 + +不涉及 + +# 四、对比分析 + +不涉及 + +# 五、设计思路与实现方案 + +## 命名与参数设计 + +本次改进不涉及新增 API，主要是对现有 `paddle.utils.cpp_extension.setup` 函数的内部实现进行改进。 + +## 底层OP设计 + +不涉及。 + +## API实现方案 + +### 1. 设置 metadata_version + +参考 https://packaging.python.org/en/latest/specifications/core-metadata/#core-metadata ， + +> Metadata consumers may want to use the more relaxed formatting rules even for metadata files that are nominally less than version 2.1. + +在 `setup` 函数中添加 `metadata_version` 参数，使用 2.1 版本的元数据格式，鼓励使用现代的元数据格式： + +```python +def setup(**attr: Any) -> None: + ... + + if 'metadata_version' not in attr: + attr['metadata_version'] = '2.1' +``` + +### 2. 扩展 BuildExtension 类 + +添加 `_generate_python_api_file` 方法，在编译完成后生成 Python API 文件： + +```python + def _generate_python_api_file(self) -> None: + """ + Generate the top-level python api file (package stub) alongside the + built shared library in build_lib. This replaces the legacy bdist_egg + write_stub mechanism that is no longer triggered in setuptools >= 80. + """ + try: + outputs = self.get_outputs() + if not outputs: + return + # We only support a single extension per setup() + so_path = os.path.abspath(outputs[0]) + so_name = os.path.basename(so_path) + build_dir = os.path.dirname(so_path) + # The package name equals distribution name + pkg_name = self.distribution.get_name() + pyfile = os.path.join(build_dir, f"{pkg_name}.py") + # Write stub; it will reference the _pd_ renamed resource at import time + custom_write_stub(so_name, pyfile) + except Exception as e: + raise RuntimeError(f"Failed to generate python api file: {e}") from e + + def run(self): + super().run() + + # Skip if using legacy bdist_egg mechanism (setuptools < 80) + if not _is_legacy_setuptools(): + # Generate python API stub into build_lib for setuptools >= 80 installs + self._generate_python_api_file() + + self._clean_intermediate_files() +``` + +### 3. 新增 InstallCommand 类 + +添加自定义的 `install` 命令类，处理以下任务： + +1. **选择合适的安装目录**：确保安装到 sys.path 中的目录 +2. **重命名共享库**：将 `{pkg}.so` 重命名为 `{pkg}_pd_.so`，避免与 Python stub 冲突 +3. **规范化包布局**：将文件组织为单一的包目录结构 + +```python +class InstallCommand(install): + """ + Extend install Command to: + 1) choose an install dir that is actually importable (on sys.path) + 2) ensure a single top-level entry for the package in site/dist-packages so + legacy tests that expect a sole artifact (egg/package) keep working + 3) rename the compiled library to *_pd_.so to avoid shadowing the python stub + """ + + def finalize_options(self) -> None: + super().finalize_options() + # Build candidate site dirs: global + user + entries already on sys.path + candidates = [] + candidates.extend(site.getsitepackages()) + usp = site.getusersitepackages() + if usp: + candidates.append(usp) + for sp in sys.path: + if isinstance(sp, str) and sp.endswith(( + 'site-packages', 'dist-packages' + )): + candidates.append(sp) + # De-dup while preserving order + seen = set() + ordered = [] + for c in candidates: + if c and c not in seen: + seen.add(c) + ordered.append(c) + # Prefer a candidate that is actually on sys.path + target = None + for c in ordered: + if c in sys.path and os.path.isdir(c): + target = c + break + # Fallback: pick the first existing candidate + if target is None: + for c in ordered: + if os.path.isdir(c): + target = c + break + if target: + self.install_lib = target + self.install_purelib = target + self.install_platlib = target + + def run(self, *args: Any, **kwargs: Any) -> None: + super().run(*args, **kwargs) + # First rename the shared library if present at top-level + self._rename_shared_library() + # Then canonicalize layout to a single top-level entry for this package + self._single_entry_layout() + + def _rename_shared_library(self) -> None: + install_dir = ( + getattr(self, 'install_lib', None) + or getattr(self, 'install_purelib', None) + or getattr(self, 'install_platlib', None) + ) + if not install_dir or not os.path.isdir(install_dir): + return + pkg = self.distribution.get_name() + suffix = ( + '.pyd' if IS_WINDOWS else ('.dylib' if OS_NAME.startswith('darwin') else '.so') + ) + old = os.path.join(install_dir, f"{pkg}{suffix}") + new = os.path.join(install_dir, f"{pkg}_pd_{suffix}") + if os.path.exists(old): + if os.path.exists(new): + os.remove(new) + os.rename(old, new) + + def _single_entry_layout(self) -> None: + """ + Ensure only one top-level item in install_dir contains the package name by: + - moving {pkg}.py -> {pkg}/__init__.py + - moving {pkg}_pd_.so -> {pkg}/{pkg}_pd_.so + - removing any {pkg}-*.egg-info left by setuptools install (only if dist-info exists) + This keeps legacy tests that scan os.listdir(site_dir) happy. + """ + install_dir = ( + getattr(self, 'install_lib', None) + or getattr(self, 'install_purelib', None) + or getattr(self, 'install_platlib', None) + ) + if not install_dir or not os.path.isdir(install_dir): + return + pkg = self.distribution.get_name() + # Check if dist-info exists + has_dist_info = any( + name.endswith('.dist-info') and name.startswith(pkg) + for name in os.listdir(install_dir) + ) + # Prepare paths + pkg_dir = os.path.join(install_dir, pkg) + py_src = os.path.join(install_dir, f"{pkg}.py") + # Find compiled lib (renamed or not) + suf_so = ( + '.pyd' if IS_WINDOWS else ('.dylib' if OS_NAME.startswith('darwin') else '.so') + ) + so_candidates = [ + os.path.join(install_dir, f"{pkg}_pd_{suf_so}"), + os.path.join(install_dir, f"{pkg}{suf_so}"), + ] + so_src = next((p for p in so_candidates if os.path.exists(p)), None) + # Create package dir + if not os.path.isdir(pkg_dir): + os.makedirs(pkg_dir, exist_ok=True) + # Move python stub to package/__init__.py if exists + if os.path.exists(py_src): + py_dst = os.path.join(pkg_dir, "__init__.py") + if os.path.exists(py_dst): + os.remove(py_dst) + os.replace(py_src, py_dst) + # Move shared lib into the package dir if exists + if so_src and os.path.exists(so_src): + so_dst = os.path.join(pkg_dir, os.path.basename(so_src)) + if os.path.exists(so_dst): + os.remove(so_dst) + os.replace(so_src, so_dst) + # Remove egg-info entries for this package only if dist-info exists + if has_dist_info: + for name in os.listdir(install_dir): + if name.startswith(f"{pkg}-") and name.endswith(".egg-info"): + p = os.path.join(install_dir, name) + if os.path.isdir(p): + shutil.rmtree(p, ignore_errors=False) + else: + os.remove(p) +``` + +### 4. 更新测试用例 + +修改测试用例中的断言，从期望 1 个包文件改为期望 2 个（egg-info + 包目录）： + +```python +# 修改前 +assert len(custom_egg_path) == 1, ( + f"Matched egg number is {len(custom_egg_path)}." +) + +# 修改后 +assert len(custom_egg_path) == 2, ( + f"Matched egg number is {len(custom_egg_path)}." +) +``` + +# 六、测试和验收的考量 + +## 测试用例 + +1. **基础功能测试**： + - 在 Setuptools 80+ 环境下编译和安装自定义算子 + - 验证生成的包结构正确 + - 验证 Python stub 文件正确生成 + - 验证共享库正确重命名 + +2. **兼容性测试**： + - 在 Setuptools >= 80 环境下验证功能正常 + - 在 Setuptools < 80 环境下验证功能正常 + +3. **导入测试**： + - 验证安装后能够正确导入自定义算子 + - 验证算子功能正常运行 + +4. **pip 集成测试**： + - 验证 `pip list` 能够正确显示已安装的自定义算子 + - 验证 `pip show` 能够正确显示自定义算子信息 + - 验证 `pip uninstall` 能够正确卸载 + +## 验收标准 + +1. 所有现有测试用例通过 +2. 在 Setuptools 80+ 环境下，自定义算子能够正常编译、安装和运行 +3. `pip list` 能够正确显示已安装的自定义算子包 +4. 不影响旧版本 Setuptools 的功能 + +# 七、可行性分析和排期规划 + +## 可行性分析 + +1. **技术可行性**：方案基于 setuptools 的标准扩展机制，技术上完全可行 +2. **兼容性风险**：通过条件判断确保兼容 80.0- 的 Setuptools +3. **测试覆盖**：现有测试用例能够覆盖主要功能点 + +## 排期规划 + +- **第 1 周**：完成核心代码实现和单元测试 +- **第 2 周**：代码审查和合并 + +# 八、影响面 + +自定义算子用户无需修改现有代码，透明升级 + +# 附件及参考资料 + +1. [Setuptools 80.0 Release Notes](https://setuptools.pypa.io/en/latest/history.html#v80-0-0) +2. [PEP 566 - Metadata for Python Software Packages 2.1](https://peps.python.org/pep-0566/) +3. [Hackathon 9th No.109 任务说明](https://github.com/PaddlePaddle/community/blob/master/hackathon/hackathon_9th/%E3%80%90Hackathon_9th%E3%80%91%E4%B8%AA%E4%BA%BA%E6%8C%91%E6%88%98%E8%B5%9B%E2%80%94%E6%A1%86%E6%9E%B6%E5%BC%80%E5%8F%91%E4%BB%BB%E5%8A%A1%E5%90%88%E9%9B%86.md#no109-基于-setuptools-80-版本自定义算子机制适配) + From e47c1712301b6a824ea6feb90590642074c30622 Mon Sep 17 00:00:00 2001 From: megemini Date: Thu, 23 Jul 2026 18:24:19 +0800 Subject: [PATCH 2/3] swe-64320 proposal --- .../PaddlePaddle__Paddle-64320/proposal.md | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/proposal.md diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/proposal.md b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/proposal.md new file mode 100644 index 000000000..fee01b5b5 --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/proposal.md @@ -0,0 +1,54 @@ +# Task Proposal: PaddlePaddle__Paddle-64320 + +## 1. 来源信息 + +- Instance ID：`PaddlePaddle__Paddle-64320` +- PR 链接：https://github.com/PaddlePaddle/Paddle/pull/64320 +- PR 标题：`【Hackathon 6th No.17】为 Paddle 新增 sparse.mask_as API -part` +- `base_commit`：`605f5e20305db0e4932a20d3e0e6cf7d7d9631d8` +- merged 时间：`2024-06-07T09:10:27Z` +- 你的身份：原 PR 作者 +- 后续联系人：megemini + +## 2. 问题一句话 + +为 Paddle 新增 `sparse.mask_as` API，支持根据给定的稀疏 mask 从稠密 Tensor 中提取对应位置的值，输出稀疏 Tensor。 + +## 3. 为什么适合作为 SWE-Paddle 样本 + +- **真实性**：该任务来自已合入的 Paddle Hackathon 6th 框架开发 PR，关联 RFC（community#901），不是合成任务。 +- **代表性**：它覆盖稀疏 CPU/GPU kernel 开发、CSR/COO 两种稀疏格式、Python API 封装、autograd 反向注册、以及 YAML op 定义，是典型的完整稀疏算子开发流程。 +- **边界清楚**：目标行为集中在 `sparse.mask_as` API 的正确实现，测试补丁可以直接暴露目标行为。CSR 格式仅支持 2-D 和 3-D，边界清晰。 +- **非平凡性**：该任务涉及 C++ kernel 实现（含 CSR 2D/3D 两种索引计算）、GPU CUDA kernel、反向梯度、Python API 封装和 YAML op 注册，不是纯 Python 或纯配置修改。 + +## 4. 任务类型和标签 + +- 任务类型：`feature_enhancement` +- 执行后端：`cpu` / `cuda` +- 设备范围：`single_gpu` +- 模块标签：`[sparse, python_api, cpu_kernel, gpu_kernel, autograd, yaml_op]` + +## 5. 验证思路 + +- 目标测试命令：`bash tests/test.sh` +- 目标测试文件：`test/legacy_test/test_sparse_mask_as_op.py` +- 修复前预期：在 `base_commit` 上应用 `tests/test.patch` 后，`test_sparse_mask_as_op.py` 中新增的 `mask_as` 相关测试应 fail（API 不存在或 kernel 未实现）。 +- 修复后预期：继续应用 `solution/code.patch` 后，目标测试应 pass。 +- P2P 候选：`test_sparse_mask_as_op.py` 为 PR 新增文件，无存量测试。建议从同模块存量稀疏测试中选取回归护栏，例如 `test_sparse_utils_op.py`、`test_sparse_unary_op.py` 等，可由 verifier 自动抽取稳定 nodeid。 + +## 6. 环境与资源 + +- 资源需求：CPU + GPU（CUDA kernel 涉及 GPU 编译） +- Paddle 来源：`PaddlePaddle/Paddle` source checkout at `base_commit` +- 是否能提供 Docker：暂无，建议后续补充 source-build Dockerfile +- patch 类型：含 C++ CPU kernel + CUDA GPU kernel + Python API + YAML op 定义 +- 环境建议：该样本涉及 C++ 和 CUDA kernel，需要 source build +- 最小测试命令：`bash tests/test.sh` +- 是否有 oracle 日志：由 SWE-Paddle verifier 结果另行维护 + +## 7. 风险自查 + +- 泄露风险：正式 `instruction.md` 应描述目标行为和验收标准，不直接指出具体修改行。 +- 环境风险：该样本涉及 C++ 和 CUDA kernel，历史 commit 复现可能需要 source build。 +- flaky 风险：需要 verifier 重复运行目标测试，并抽取稳定 F2P/P2P nodeid。 +- 拆分风险：该 PR 的目标集中在新增 `sparse.mask_as` API，适合作为一个样本。 From b2a3263410dc1debc8cb3f4e7337c5df491efb73 Mon Sep 17 00:00:00 2001 From: megemini Date: Fri, 24 Jul 2026 14:08:10 +0800 Subject: [PATCH 3/3] swe paddle 64320 codes --- .../PaddlePaddle__Paddle-64320/README.md | 37 + .../environment/README.md | 27 + .../PaddlePaddle__Paddle-64320/instruction.md | 24 + .../solution/code.patch | 872 ++++++++++++++++++ .../tests/test.patch | 176 ++++ .../PaddlePaddle__Paddle-64320/tests/test.sh | 4 + 6 files changed, 1140 insertions(+) create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/README.md create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/environment/README.md create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/instruction.md create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/solution/code.patch create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.patch create mode 100644 swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.sh diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/README.md b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/README.md new file mode 100644 index 000000000..a35cd7bbf --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/README.md @@ -0,0 +1,37 @@ +# PaddlePaddle__Paddle-64320 + +This directory converts Paddle PR #64320 into a SWE-Paddle community task candidate. + +## Source + +| Field | Value | +| --- | --- | +| Repo | `PaddlePaddle/Paddle` | +| PR | [64320](https://github.com/PaddlePaddle/Paddle/pull/64320) | +| PR title | 【Hackathon 6th No.17】为 Paddle 新增 sparse.mask_as API -part | +| Base commit | `605f5e20305db0e4932a20d3e0e6cf7d7d9631d8` | +| Merged at | `2024-06-07T09:10:27Z` | +| Hackathon | `6th` task `17` | +| Task type | `feature_enhancement` | +| Resource | CPU + GPU | + +## Summary + +Add `sparse.mask_as` API for Paddle, which extracts values from a dense tensor at positions indicated by a sparse mask and outputs a sparse tensor. + +## Files + +- `proposal.md`: candidate proposal for maintainer triage. +- `instruction.md`: self-contained problem statement for the coding agent. +- `solution/code.patch`: gold patch from the merged PR. +- `tests/test.patch`: test patch exposing the target behavior. +- `tests/test.sh`: minimal target test command. +- `environment/README.md`: environment notes for reproduction. + +## Verification + +```bash +bash tests/test.sh +``` + +Expected behavior: applying `tests/test.patch` to `base_commit` should fail on the target behavior; applying both `tests/test.patch` and `solution/code.patch` should pass the target tests. diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/environment/README.md b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/environment/README.md new file mode 100644 index 000000000..85bd29e04 --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/environment/README.md @@ -0,0 +1,27 @@ +# Environment Notes + +This candidate is part of the SWE-Paddle task set. + +## Expected Environment + +- Repository: `PaddlePaddle/Paddle` +- Base commit: `605f5e20305db0e4932a20d3e0e6cf7d7d9631d8` +- Resource: CPU + GPU +- GPU required: yes (CUDA kernels are included) +- Build path: Paddle source checkout at the base commit. This task involves C++ CPU kernels, CUDA GPU kernels, Python API, and YAML op definitions, so source build is required. + +## Run Order + +1. Check out `PaddlePaddle/Paddle` at the base commit. +2. Apply `tests/test.patch`. +3. Run `bash tests/test.sh`; the target behavior should fail before the fix. +4. Apply `solution/code.patch`. +5. Run `bash tests/test.sh` again; the target behavior should pass after the gold patch. + +## Minimal Test Command + +```bash +bash tests/test.sh +``` + +The verifier is responsible for deriving stable F2P and P2P node IDs from repeated runs. diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/instruction.md b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/instruction.md new file mode 100644 index 000000000..6238ee3ad --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/instruction.md @@ -0,0 +1,24 @@ +# 新增 sparse.mask_as API + +## 详细描述 + +为 Paddle 稀疏计算新增 `sparse.mask_as` API。该 API 根据给定的稀疏 mask（SparseCooTensor 或 SparseCsrTensor），从稠密 Tensor 中提取对应非零位置的值，输出一个与 mask 具有相同 indices 的稀疏 Tensor。 + +要求支持： +- COO 格式：支持 1-D ~ 4-D 输入 +- CSR 格式：支持 2-D 和 3-D 输入（其他维度应报错） +- 数据类型：float32, float64, int32, int64, complex64, complex128, int8, int16, float16 +- 前向计算 + 反向梯度 + +## 验收说明 + +- `paddle.sparse.mask_as(x, mask)` 应能正确根据 mask 的 indices 从稠密 Tensor x 中提取值 +- 支持 COO 和 CSR 两种稀疏格式 +- CSR 格式仅支持 2-D 和 3-D，其他维度应报错 +- 反向梯度应正确传播 + +## Acceptance Criteria + +- The behavior described above should be fixed. +- Existing valid behavior should remain unchanged. +- Do not satisfy the task by deleting tests, weakening assertions, or bypassing validation broadly. diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/solution/code.patch b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/solution/code.patch new file mode 100644 index 000000000..602a3c2b4 --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/solution/code.patch @@ -0,0 +1,872 @@ +diff --git a/paddle/phi/kernels/sparse/cpu/mask_grad_kernel.cc b/paddle/phi/kernels/sparse/cpu/mask_grad_kernel.cc +new file mode 100644 +index 00000000000000..3503c88b2ef8b4 +--- /dev/null ++++ b/paddle/phi/kernels/sparse/cpu/mask_grad_kernel.cc +@@ -0,0 +1,56 @@ ++/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. ++ ++Licensed under the Apache License, Version 2.0 (the "License"); ++you may not use this file except in compliance with the License. ++You may obtain a copy of the License at ++ ++ http://www.apache.org/licenses/LICENSE-2.0 ++ ++Unless required by applicable law or agreed to in writing, software ++distributed under the License is distributed on an "AS IS" BASIS, ++WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++See the License for the specific language governing permissions and ++limitations under the License. */ ++ ++#include "paddle/phi/kernels/sparse/mask_grad_kernel.h" ++#include "paddle/phi/kernels/sparse/mask_kernel.h" ++#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h" ++ ++#include "paddle/phi/backends/cpu/cpu_context.h" ++#include "paddle/phi/core/kernel_registry.h" ++ ++PD_REGISTER_KERNEL(mask_as_coo_grad, ++ CPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskAsCooGradKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int8_t, ++ int16_t, ++ int, ++ int64_t, ++ bool, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); ++} ++ ++PD_REGISTER_KERNEL(mask_as_csr_grad, ++ CPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskAsCsrGradKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int8_t, ++ int16_t, ++ int, ++ int64_t, ++ bool, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR); ++} +diff --git a/paddle/phi/kernels/sparse/cpu/mask_kernel.cc b/paddle/phi/kernels/sparse/cpu/mask_kernel.cc +index 5213dd44a4c07c..658a26452dafb6 100644 +--- a/paddle/phi/kernels/sparse/cpu/mask_kernel.cc ++++ b/paddle/phi/kernels/sparse/cpu/mask_kernel.cc +@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and + limitations under the License. */ + + #include "paddle/phi/kernels/sparse/mask_kernel.h" ++#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h" + + #include "paddle/common/ddim.h" + #include "paddle/phi/api/ext/dispatch.h" +@@ -75,16 +76,116 @@ void MaskCooCPUKernel(const CPUContext& dev_ctx, + * x and mask must have the same shape. + **/ + template +-void MaskCooKernel(const Context& dev_ctx, +- const DenseTensor& x, +- const SparseCooTensor& mask, +- SparseCooTensor* out) { ++void MaskAsCooKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCooTensor& mask, ++ SparseCooTensor* out) { + PD_VISIT_BASE_INTEGRAL_TYPES( + mask.indices().dtype(), "MaskCooCPUKernel", ([&] { + MaskCooCPUKernel(dev_ctx, x, mask, out); + })); + } + ++template ++void MaskCsr2DCPUKernel(const CPUContext& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const DenseTensor& mask_cols = mask.cols(); ++ const DenseTensor& mask_crows = mask.crows(); ++ int64_t num_non_zeros = mask.nnz(); ++ ++ DenseTensor out_cols = phi::EmptyLike(dev_ctx, mask_cols); ++ DenseTensor out_crows = phi::EmptyLike(dev_ctx, mask_crows); ++ DenseTensor out_values = phi::Empty(dev_ctx, {num_non_zeros}); ++ ++ phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols); ++ phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows); ++ ++ int64_t numel = 0; ++ for (int64_t i = 0; i < mask_crows.numel() - 1; ++i) { ++ for (int64_t j = mask_crows.data()[i]; ++ j < mask_crows.data()[i + 1]; ++ ++j) { ++ IntT col_idx = mask_cols.data()[numel]; ++ ++ out_values.data()[numel] = ++ x.data()[(i / x.dims()[0]) * x.dims()[1] + ++ (i % x.dims()[0]) * x.dims()[1] + col_idx]; ++ ++ ++numel; ++ } ++ } ++ ++ out->SetMember(out_crows, out_cols, out_values, x.dims()); ++} ++ ++template ++void MaskCsr3DCPUKernel(const CPUContext& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const DenseTensor& mask_cols = mask.cols(); ++ const DenseTensor& mask_crows = mask.crows(); ++ int64_t num_non_zeros = mask.nnz(); ++ ++ DenseTensor out_cols = phi::EmptyLike(dev_ctx, mask_cols); ++ DenseTensor out_crows = phi::EmptyLike(dev_ctx, mask_crows); ++ DenseTensor out_values = phi::Empty(dev_ctx, {num_non_zeros}); ++ ++ phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols); ++ phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows); ++ ++ int64_t numel = 0; ++ for (int64_t i = 0; i < mask_crows.numel() - 1; ++i) { ++ for (int64_t j = mask_crows.data()[i]; ++ j < mask_crows.data()[i + 1]; ++ ++j) { ++ IntT col_idx = mask_cols.data()[numel]; ++ ++ out_values.data()[numel] = ++ x.data()[(i / (mask_crows.numel() / x.dims()[0])) * ++ (x.dims()[1] * x.dims()[2]) + ++ (i % (mask_crows.numel() / x.dims()[0])) * x.dims()[2] + ++ col_idx]; ++ ++ ++numel; ++ } ++ } ++ ++ out->SetMember(out_crows, out_cols, out_values, x.dims()); ++} ++ ++/** ++ * @brief Filter the DenseTensor x by the ++ * mask.crows(), mask.cols() and output a SparseCsrTensor ++ * x and mask must have the same shape. ++ **/ ++template ++void MaskAsCsrKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const phi::DDim& x_dims = x.dims(); ++ if (x_dims.size() == 2) { ++ PD_VISIT_BASE_INTEGRAL_TYPES( ++ mask.crows().dtype(), "MaskCsr2DCPUKernel", ([&] { ++ MaskCsr2DCPUKernel(dev_ctx, x, mask, out); ++ })); ++ } else if (x_dims.size() == 3) { ++ PD_VISIT_BASE_INTEGRAL_TYPES( ++ mask.crows().dtype(), "MaskCsr3DCPUKernel", ([&] { ++ MaskCsr3DCPUKernel(dev_ctx, x, mask, out); ++ })); ++ } else { ++ // throw exception ++ phi::errors::InvalidArgument( ++ "mask_as for Sparse CSR Tensor only support 2-D or 3-D, but got " ++ "%d-D.", ++ x_dims.size()); ++ } ++} ++ + template + void MaskHelperCooCPUKernel(const CPUContext& dev_ctx, + const SparseCooTensor& x, +@@ -157,10 +258,26 @@ void MaskHelperCooKernel(const Context& dev_ctx, + } // namespace sparse + } // namespace phi + +-PD_REGISTER_KERNEL(mask_coo, ++PD_REGISTER_KERNEL(mask_helper_coo, ++ CPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskHelperCooKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int16_t, ++ int, ++ int64_t, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); ++} ++ ++PD_REGISTER_KERNEL(mask_as_coo, + CPU, + ALL_LAYOUT, +- phi::sparse::MaskCooKernel, ++ phi::sparse::MaskAsCooKernel, + float, + double, + uint8_t, +@@ -174,18 +291,19 @@ PD_REGISTER_KERNEL(mask_coo, + kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); + } + +-PD_REGISTER_KERNEL(mask_helper_coo, ++PD_REGISTER_KERNEL(mask_as_csr, + CPU, + ALL_LAYOUT, +- phi::sparse::MaskHelperCooKernel, ++ phi::sparse::MaskAsCsrKernel, + float, + double, +- phi::dtype::float16, + uint8_t, ++ int8_t, + int16_t, + int, + int64_t, ++ bool, + phi::dtype::complex, + phi::dtype::complex) { +- kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR); + } +diff --git a/paddle/phi/kernels/sparse/gpu/mask_grad_kernel.cu b/paddle/phi/kernels/sparse/gpu/mask_grad_kernel.cu +new file mode 100644 +index 00000000000000..1e4e3276d82e15 +--- /dev/null ++++ b/paddle/phi/kernels/sparse/gpu/mask_grad_kernel.cu +@@ -0,0 +1,56 @@ ++/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. ++ ++Licensed under the Apache License, Version 2.0 (the "License"); ++you may not use this file except in compliance with the License. ++You may obtain a copy of the License at ++ ++ http://www.apache.org/licenses/LICENSE-2.0 ++ ++Unless required by applicable law or agreed to in writing, software ++distributed under the License is distributed on an "AS IS" BASIS, ++WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++See the License for the specific language governing permissions and ++limitations under the License. */ ++ ++#include "paddle/phi/kernels/sparse/mask_grad_kernel.h" ++#include "paddle/phi/kernels/sparse/mask_kernel.h" ++#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h" ++ ++#include "paddle/phi/backends/cpu/cpu_context.h" ++#include "paddle/phi/core/kernel_registry.h" ++ ++PD_REGISTER_KERNEL(mask_as_coo_grad, ++ GPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskAsCooGradKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int8_t, ++ int16_t, ++ int, ++ int64_t, ++ bool, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); ++} ++ ++PD_REGISTER_KERNEL(mask_as_csr_grad, ++ GPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskAsCsrGradKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int8_t, ++ int16_t, ++ int, ++ int64_t, ++ bool, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR); ++} +diff --git a/paddle/phi/kernels/sparse/gpu/mask_kernel.cu b/paddle/phi/kernels/sparse/gpu/mask_kernel.cu +index 0941ad69b0dd2d..3459f6802b8819 100644 +--- a/paddle/phi/kernels/sparse/gpu/mask_kernel.cu ++++ b/paddle/phi/kernels/sparse/gpu/mask_kernel.cu +@@ -12,7 +12,10 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. */ + ++#include ++ + #include "paddle/phi/kernels/sparse/mask_kernel.h" ++#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h" + + #include "paddle/common/ddim.h" + #include "paddle/phi/backends/gpu/gpu_info.h" +@@ -106,22 +109,256 @@ void MaskCooGPUKernel(const GPUContext& dev_ctx, + out->SetMember(out_indices, out_values, dims, true); + } + ++template ++__global__ void ConvertCsrCrowsToCooRows(const IntT* crows_ptr, ++ const IntT* crows_offsets, ++ IntT* rows_ptr, ++ IntT* batch_ptr, ++ const int rows) { ++ const int b = blockIdx.y; ++ const int64_t offset = crows_offsets ? crows_offsets[b] : 0; ++ const int tid = threadIdx.x + blockIdx.x * blockDim.x; ++ for (int i = tid; i < rows; i += gridDim.x * blockDim.x) { ++ for (int j = crows_ptr[b * (rows + 1) + i]; ++ j < crows_ptr[b * (rows + 1) + i + 1]; ++ j++) { ++ rows_ptr[offset + j] = i; ++ if (batch_ptr) { ++ batch_ptr[offset + j] = b; ++ } ++ } ++ } ++} ++ ++template ++__global__ void GetBatchSizes(const IntT* crows, ++ const int rows, ++ const int batches, ++ IntT* batch_sizes) { ++ const int tid = threadIdx.x + blockIdx.x * blockDim.x; ++ if (tid < batches) { ++ batch_sizes[tid] = crows[tid * (rows + 1) + rows]; ++ } ++} ++ ++template ++void MaskCsr2DGPUKernel(const GPUContext& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const DenseTensor& mask_cols = mask.cols(); ++ const DenseTensor& mask_crows = mask.crows(); ++ int64_t num_non_zeros = mask.nnz(); ++ ++ DenseTensor out_cols = phi::EmptyLike(dev_ctx, mask_cols); ++ DenseTensor out_crows = phi::EmptyLike(dev_ctx, mask_crows); ++ DenseTensor out_values = phi::Empty(dev_ctx, {num_non_zeros}); ++ ++ phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols); ++ phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows); ++ ++ const DDim& dims = x.dims(); ++ const int64_t non_zero_num = mask.nnz(); ++ int64_t sparse_dim = 2; ++ DenseTensor sparse_offsets = phi::Empty(dev_ctx, {sparse_dim}); ++ std::vector h_sparse_offsets(sparse_dim); ++ phi::funcs::sparse::CalcOffsetsPerDim( ++ dims, sparse_dim, h_sparse_offsets.data()); ++ ++ phi::backends::gpu::GpuMemcpyAsync(sparse_offsets.data(), ++ &h_sparse_offsets[0], ++ sizeof(int64_t) * sparse_dim, ++ gpuMemcpyHostToDevice, ++ dev_ctx.stream()); ++ ++ const auto& csr_crows = mask.crows(); ++ const auto& csr_cols = mask.cols(); ++ const IntT* csr_crows_data = csr_crows.data(); ++ const IntT* csr_cols_data = csr_cols.data(); ++ ++ const int batches = 1; ++ const int rows = dims[0]; ++ auto dims_2d = flatten_to_2d(dims, sparse_dim); ++ const int cols = dims_2d[1]; ++ ++ DenseTensor indices = phi::Empty(dev_ctx, {sparse_dim, non_zero_num}); ++ IntT* coo_indices = indices.data(); ++ IntT* batch_ptr = nullptr; ++ IntT* coo_rows_data = coo_indices; ++ IntT* coo_cols_data = coo_rows_data + non_zero_num; ++ IntT* offsets_ptr = nullptr; ++ ++ auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1); ++ config.block_per_grid.y = batches; ++ ConvertCsrCrowsToCooRows ++ <<>>( ++ csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows); ++ phi::backends::gpu::GpuMemcpyAsync(coo_cols_data, ++ csr_cols_data, ++ sizeof(IntT) * non_zero_num, ++ gpuMemcpyDeviceToDevice, ++ dev_ctx.stream()); ++ ++ const T* x_ptr = x.data(); ++ const IntT* indices_ptr = coo_indices; ++ T* out_values_ptr = out_values.data(); ++ ++ auto config_mask = ++ phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1); ++ MaskKernel<<>>(x_ptr, ++ indices_ptr, ++ sparse_offsets.data(), ++ non_zero_num, ++ cols, ++ sparse_dim, ++ out_values_ptr); ++ ++ out->SetMember(out_crows, out_cols, out_values, x.dims()); ++} ++ ++template ++void MaskCsr3DGPUKernel(const GPUContext& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const DenseTensor& mask_cols = mask.cols(); ++ const DenseTensor& mask_crows = mask.crows(); ++ int64_t num_non_zeros = mask.nnz(); ++ ++ DenseTensor out_cols = phi::EmptyLike(dev_ctx, mask_cols); ++ DenseTensor out_crows = phi::EmptyLike(dev_ctx, mask_crows); ++ DenseTensor out_values = phi::Empty(dev_ctx, {num_non_zeros}); ++ ++ phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols); ++ phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows); ++ ++ const DDim& dims = x.dims(); ++ const int64_t non_zero_num = mask.nnz(); ++ int64_t sparse_dim = 3; ++ DenseTensor sparse_offsets = phi::Empty(dev_ctx, {sparse_dim}); ++ std::vector h_sparse_offsets(sparse_dim); ++ phi::funcs::sparse::CalcOffsetsPerDim( ++ dims, sparse_dim, h_sparse_offsets.data()); ++ ++ phi::backends::gpu::GpuMemcpyAsync(sparse_offsets.data(), ++ &h_sparse_offsets[0], ++ sizeof(int64_t) * sparse_dim, ++ gpuMemcpyHostToDevice, ++ dev_ctx.stream()); ++ ++ const auto& csr_crows = mask.crows(); ++ const auto& csr_cols = mask.cols(); ++ const IntT* csr_crows_data = csr_crows.data(); ++ const IntT* csr_cols_data = csr_cols.data(); ++ ++ const int batches = dims[0]; ++ const int rows = dims[1]; ++ auto dims_2d = flatten_to_2d(dims, sparse_dim); ++ const int cols = dims_2d[1]; ++ ++ DenseTensor indices = phi::Empty(dev_ctx, {sparse_dim, non_zero_num}); ++ DenseTensor offsets = phi::Empty(dev_ctx, {batches}); ++ IntT* coo_indices = indices.data(); ++ IntT* batch_ptr = coo_indices; ++ IntT* coo_rows_data = batch_ptr + non_zero_num; ++ IntT* coo_cols_data = coo_rows_data + non_zero_num; ++ IntT* offsets_ptr = offsets.data(); ++ ++ auto config_batch = ++ phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, batches, 1); ++ GetBatchSizes ++ <<>>( ++ csr_crows_data, rows, batches, offsets_ptr); ++ ++#ifdef PADDLE_WITH_HIP ++ thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()), ++#else ++ thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()), ++#endif ++ offsets_ptr, ++ offsets_ptr + batches, ++ offsets_ptr); ++ ++ auto config = phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1); ++ config.block_per_grid.y = batches; ++ ConvertCsrCrowsToCooRows ++ <<>>( ++ csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows); ++ phi::backends::gpu::GpuMemcpyAsync(coo_cols_data, ++ csr_cols_data, ++ sizeof(IntT) * non_zero_num, ++ gpuMemcpyDeviceToDevice, ++ dev_ctx.stream()); ++ ++ const T* x_ptr = x.data(); ++ const IntT* indices_ptr = coo_indices; ++ T* out_values_ptr = out_values.data(); ++ ++ auto config_mask = ++ phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1); ++ MaskKernel<<>>(x_ptr, ++ indices_ptr, ++ sparse_offsets.data(), ++ non_zero_num, ++ cols, ++ sparse_dim, ++ out_values_ptr); ++ ++ out->SetMember(out_crows, out_cols, out_values, x.dims()); ++} ++ + /** + * @brief Filter the DenseTensor x by the + * mask.indices() and output a SparseCooTensor + * x and mask must have the same shape. + **/ + template +-void MaskCooKernel(const Context& dev_ctx, +- const DenseTensor& x, +- const SparseCooTensor& mask, +- SparseCooTensor* out) { ++void MaskAsCooKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCooTensor& mask, ++ SparseCooTensor* out) { + PD_VISIT_BASE_INTEGRAL_TYPES( + mask.indices().dtype(), "MaskCooGPUKernel", ([&] { + MaskCooGPUKernel(dev_ctx, x, mask, out); + })); + } + ++/** ++ * @brief Filter the DenseTensor x by the ++ * mask.crows(), mask.cols() and output a SparseCsrTensor ++ * x and mask must have the same shape. ++ **/ ++template ++void MaskAsCsrKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out) { ++ const phi::DDim& x_dims = x.dims(); ++ if (x_dims.size() == 2) { ++ PD_VISIT_BASE_INTEGRAL_TYPES( ++ mask.crows().dtype(), "MaskCsr2DGPUKernel", ([&] { ++ MaskCsr2DGPUKernel(dev_ctx, x, mask, out); ++ })); ++ } else if (x_dims.size() == 3) { ++ PD_VISIT_BASE_INTEGRAL_TYPES( ++ mask.crows().dtype(), "MaskCsr3DGPUKernel", ([&] { ++ MaskCsr3DGPUKernel(dev_ctx, x, mask, out); ++ })); ++ } else { ++ // throw exception ++ phi::errors::InvalidArgument( ++ "mask_as for Sparse CSR Tensor only support 2-D or 3-D, but got " ++ "%d-D.", ++ x_dims.size()); ++ } ++} ++ + template + __global__ void MaskTable(const IntT* x_indexs, + const int n, +@@ -296,10 +533,26 @@ void MaskHelperCooKernel(const Context& dev_ctx, + } // namespace sparse + } // namespace phi + +-PD_REGISTER_KERNEL(mask_coo, ++PD_REGISTER_KERNEL(mask_helper_coo, + GPU, + ALL_LAYOUT, +- phi::sparse::MaskCooKernel, ++ phi::sparse::MaskHelperCooKernel, ++ float, ++ double, ++ phi::dtype::float16, ++ uint8_t, ++ int16_t, ++ int, ++ int64_t, ++ phi::dtype::complex, ++ phi::dtype::complex) { ++ kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); ++} ++ ++PD_REGISTER_KERNEL(mask_as_coo, ++ GPU, ++ ALL_LAYOUT, ++ phi::sparse::MaskAsCooKernel, + float, + double, + phi::dtype::float16, +@@ -314,18 +567,20 @@ PD_REGISTER_KERNEL(mask_coo, + kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); + } + +-PD_REGISTER_KERNEL(mask_helper_coo, ++PD_REGISTER_KERNEL(mask_as_csr, + GPU, + ALL_LAYOUT, +- phi::sparse::MaskHelperCooKernel, ++ phi::sparse::MaskAsCsrKernel, + float, + double, + phi::dtype::float16, + uint8_t, ++ int8_t, + int16_t, + int, + int64_t, ++ bool, + phi::dtype::complex, + phi::dtype::complex) { +- kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); ++ kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR); + } +diff --git a/paddle/phi/kernels/sparse/mask_grad_kernel.h b/paddle/phi/kernels/sparse/mask_grad_kernel.h +new file mode 100644 +index 00000000000000..687562aa300d1c +--- /dev/null ++++ b/paddle/phi/kernels/sparse/mask_grad_kernel.h +@@ -0,0 +1,45 @@ ++/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. ++ ++Licensed under the Apache License, Version 2.0 (the "License"); ++you may not use this file except in compliance with the License. ++You may obtain a copy of the License at ++ ++ http://www.apache.org/licenses/LICENSE-2.0 ++ ++Unless required by applicable law or agreed to in writing, software ++distributed under the License is distributed on an "AS IS" BASIS, ++WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++See the License for the specific language governing permissions and ++limitations under the License. */ ++ ++#pragma once ++ ++#include "paddle/phi/core/dense_tensor.h" ++#include "paddle/phi/core/sparse_coo_tensor.h" ++#include "paddle/phi/core/sparse_csr_tensor.h" ++#include "paddle/phi/kernels/sparse/mask_kernel.h" ++#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h" ++ ++namespace phi { ++namespace sparse { ++ ++template ++void MaskAsCooGradKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCooTensor& mask, ++ const SparseCooTensor& out_grad, ++ DenseTensor* x_grad) { ++ CooToDenseKernel(dev_ctx, out_grad, x_grad); ++} ++ ++template ++void MaskAsCsrGradKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ const SparseCsrTensor& out_grad, ++ DenseTensor* x_grad) { ++ CsrToDenseKernel(dev_ctx, out_grad, x_grad); ++} ++ ++} // namespace sparse ++} // namespace phi +diff --git a/paddle/phi/kernels/sparse/mask_kernel.h b/paddle/phi/kernels/sparse/mask_kernel.h +index 5ffc7fb4aa44d9..5be993e243b193 100644 +--- a/paddle/phi/kernels/sparse/mask_kernel.h ++++ b/paddle/phi/kernels/sparse/mask_kernel.h +@@ -16,21 +16,28 @@ limitations under the License. */ + + #include "paddle/phi/core/dense_tensor.h" + #include "paddle/phi/core/sparse_coo_tensor.h" ++#include "paddle/phi/core/sparse_csr_tensor.h" + + namespace phi { + namespace sparse { + +-template +-void MaskCooKernel(const Context& dev_ctx, +- const DenseTensor& x, +- const SparseCooTensor& mask, +- SparseCooTensor* out); +- + template + void MaskHelperCooKernel(const Context& dev_ctx, + const SparseCooTensor& x, + const DenseTensor& mask_indices, + DenseTensor* out); + ++template ++void MaskAsCooKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCooTensor& mask, ++ SparseCooTensor* out); ++ ++template ++void MaskAsCsrKernel(const Context& dev_ctx, ++ const DenseTensor& x, ++ const SparseCsrTensor& mask, ++ SparseCsrTensor* out); ++ + } // namespace sparse + } // namespace phi +diff --git a/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc b/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc +index f5915c7acb84ce..2b802615486f42 100644 +--- a/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc ++++ b/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc +@@ -32,7 +32,7 @@ void CooToDenseGradKernel(const Context& dev_ctx, + const SparseCooTensor& x, + const DenseTensor& out_grad, + SparseCooTensor* x_grad) { +- MaskCooKernel(dev_ctx, out_grad, x, x_grad); ++ MaskAsCooKernel(dev_ctx, out_grad, x, x_grad); + } + + } // namespace sparse +diff --git a/paddle/phi/ops/yaml/sparse_backward.yaml b/paddle/phi/ops/yaml/sparse_backward.yaml +index 3e614b942d3019..f7734af1bf6ecd 100644 +--- a/paddle/phi/ops/yaml/sparse_backward.yaml ++++ b/paddle/phi/ops/yaml/sparse_backward.yaml +@@ -184,6 +184,17 @@ + func : log1p_coo_grad {sparse_coo, sparse_coo -> sparse_coo}, + log1p_csr_grad {sparse_csr, sparse_csr -> sparse_csr} + ++- backward_op : mask_as_grad ++ forward : mask_as(Tensor x, Tensor mask) -> Tensor(out) ++ args : (Tensor x, Tensor mask, Tensor out_grad) ++ output : Tensor(x_grad) ++ infer_meta : ++ func : UnchangedInferMeta ++ param : [x] ++ kernel : ++ func : mask_as_coo_grad {dense, sparse_coo, sparse_coo -> dense}, ++ mask_as_csr_grad {dense, sparse_csr, sparse_csr -> dense} ++ + - backward_op : masked_matmul_grad + forward : masked_matmul(Tensor x, Tensor y, Tensor mask) -> Tensor(out) + args : (Tensor x, Tensor y, Tensor out_grad) +diff --git a/paddle/phi/ops/yaml/sparse_ops.yaml b/paddle/phi/ops/yaml/sparse_ops.yaml +index ac230be485c095..80cef73a6c1f5e 100644 +--- a/paddle/phi/ops/yaml/sparse_ops.yaml ++++ b/paddle/phi/ops/yaml/sparse_ops.yaml +@@ -497,6 +497,18 @@ + func : indices_coo{sparse_coo -> dense} + layout : x + ++- op: mask_as ++ args : (Tensor x, Tensor mask) ++ output : Tensor(out) ++ infer_meta : ++ func : UnchangedInferMeta ++ param : [x] ++ kernel : ++ func : mask_as_coo{dense, sparse_coo -> sparse_coo}, ++ mask_as_csr{dense, sparse_csr -> sparse_csr} ++ layout : x ++ backward: mask_as_grad ++ + - op: masked_matmul + args : (Tensor x, Tensor y, Tensor mask) + output : Tensor(out) +diff --git a/python/paddle/sparse/__init__.py b/python/paddle/sparse/__init__.py +index 661143f12dae8f..98f5ca0b13ee54 100644 +--- a/python/paddle/sparse/__init__.py ++++ b/python/paddle/sparse/__init__.py +@@ -17,6 +17,7 @@ + add, + divide, + is_same_shape, ++ mask_as, + masked_matmul, + matmul, + multiply, +@@ -77,6 +78,7 @@ + 'expm1', + 'mv', + 'matmul', ++ 'mask_as', + 'masked_matmul', + 'addmm', + 'add', +diff --git a/python/paddle/sparse/binary.py b/python/paddle/sparse/binary.py +index 3aac3d5e7f1446..abc943ac3c1fc2 100644 +--- a/python/paddle/sparse/binary.py ++++ b/python/paddle/sparse/binary.py +@@ -452,3 +452,60 @@ def is_same_shape(x, y): + + """ + return x.is_same_shape(y) ++ ++ ++@dygraph_only ++def mask_as(x, mask, name=None): ++ r""" ++ Filter the input dense tensor `x` using the `indices` of the sparse matrix `mask`, ++ which in turn generates a sparse matrix of the corresponding format. ++ The input `x` and `mask` must have the same shape, and the sparse tensor returned has the same indices as `mask` ++ even `zero` values exist in the coresponding indices. ++ ++ Args: ++ x (Tensor): The input tensor. It should be a DenseTensor. ++ The data type can be float32, float64, int32, int64, complex64, complex128, int8, int16, float16. ++ mask (Tensor): The input tensor. It can be SparseCooTensor or SparseCsrTensor. ++ It should be 2D or 3D when the mask is SparseCsrTensor. ++ name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`. ++ ++ Returns: ++ Tensor: A sparse tensor. ++ ++ Examples: ++ .. code-block:: python ++ ++ >>> import paddle ++ >>> paddle.set_device('cpu') ++ ++ >>> # csr sparse tensor ++ >>> crows = [0, 2, 3, 5] ++ >>> cols = [1, 3, 2, 0, 1] ++ >>> values = [1., 2., 3., 4., 5.] ++ >>> dense_shape = [3, 4] ++ >>> csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape) ++ >>> paddle.seed(2024) ++ >>> x = paddle.rand(dense_shape).astype(csr.dtype) ++ >>> out = paddle.sparse.mask_as(x, csr) ++ >>> print(out) ++ Tensor(shape=[3, 4], dtype=paddle.float32, place=Place(cpu), stop_gradient=True, ++ crows=[0, 2, 3, 5], ++ cols=[1, 3, 2, 0, 1], ++ values=[0.23659813, 0.08467803, 0.64152628, 0.66596609, 0.90394485]) ++ ++ >>> # coo sparse tensor ++ >>> indices = [[0, 1, 2], [1, 2, 0]] ++ >>> values = [1.0, 2.0, 3.0] ++ >>> dense_shape = [3, 3] ++ >>> coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape) ++ >>> paddle.seed(2024) ++ >>> x = paddle.rand(dense_shape).astype(coo.dtype) ++ >>> out = paddle.sparse.mask_as(x, coo) ++ >>> print(out) ++ Tensor(shape=[3, 3], dtype=paddle.float32, place=Place(cpu), stop_gradient=True, ++ indices=[[0, 1, 2], ++ [1, 2, 0]], ++ values=[0.23659813, 0.40340215, 0.64152628]) ++ ++ """ ++ return _C_ops.sparse_mask_as(x, mask) diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.patch b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.patch new file mode 100644 index 000000000..9cbfe9c34 --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.patch @@ -0,0 +1,176 @@ +diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt +index f84458dd494f30..33aa88c3a7c516 100644 +--- a/test/legacy_test/CMakeLists.txt ++++ b/test/legacy_test/CMakeLists.txt +@@ -1131,6 +1131,7 @@ set_pir_tests_properties() + set_tests_properties(test_nadam_op PROPERTIES TIMEOUT 100) + set_tests_properties(test_radam_op PROPERTIES TIMEOUT 100) + set_tests_properties(test_nan_inf PROPERTIES TIMEOUT 120) ++set_tests_properties(test_sparse_mask_as_op PROPERTIES TIMEOUT 120) + set_tests_properties(test_bicubic_interp_op PROPERTIES TIMEOUT 120) + set_tests_properties(test_bilinear_interp_op PROPERTIES TIMEOUT 120) + set_tests_properties(test_conv2d_op_depthwise_conv +diff --git a/test/legacy_test/test_sparse_mask_as_op.py b/test/legacy_test/test_sparse_mask_as_op.py +new file mode 100644 +index 00000000000000..f4cd639452b5d3 +--- /dev/null ++++ b/test/legacy_test/test_sparse_mask_as_op.py +@@ -0,0 +1,159 @@ ++# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++import unittest ++ ++import numpy as np ++ ++import paddle ++ ++ ++def generate_data(shape, dtype): ++ """ ++ Generate `data` and `mask` with the same shape and dtype. ++ """ ++ _mask = np.random.randint(0, 2, shape) ++ if np.sum(_mask) == 0: ++ _mask.flat[0] = 1 ++ mask = (np.random.randint(-100, 100, shape) * _mask).astype(dtype) ++ data = np.random.randint(-100, 100, shape).astype(dtype) ++ return data, mask ++ ++ ++class TestMaskAs(unittest.TestCase): ++ def setUp(self): ++ self.init_format() ++ self.places = [paddle.CPUPlace()] ++ if paddle.is_compiled_with_cuda(): ++ self.places.append(paddle.CUDAPlace(0)) ++ ++ def init_format(self): ++ self.format = None ++ ++ def check(self, shape, dtype, place, check_grad=True): ++ paddle.disable_static() ++ dense_data_np, dense_mask_np = generate_data(shape, dtype) ++ ++ dense_data_pd = paddle.to_tensor( ++ dense_data_np, dtype=dtype, place=place ++ ) ++ dense_data_pd.stop_gradient = False ++ ++ if self.format == 'coo': ++ sparse_mask_pd = paddle.to_tensor( ++ dense_mask_np, dtype=dtype, place=place ++ ).to_sparse_coo(len(shape)) ++ else: ++ sparse_mask_pd = paddle.to_tensor( ++ dense_mask_np, dtype=dtype, place=place ++ ).to_sparse_csr() ++ ++ sparse_out_pd = paddle.sparse.mask_as(dense_data_pd, sparse_mask_pd) ++ ++ # compare the tensor from sparse->dense with reference numpy data ++ # the result only keeps the values where mask not zero, like: ++ # dense_data_np ++ # [[ 38. 15. 76.] ++ # [-98. -75. 10.] ++ # [-52. 49. -48.]] ++ # dense_mask_np ++ # [[-70. 0. 0.] ++ # [-50. 34. 60.] ++ # [-34. 0. -18.]] ++ # dense_data_np_ref ++ # [[ 38. 0. 0.] ++ # [-98. -75. 10.] ++ # [-52. 0. -48.]] ++ dense_data_np_ref = dense_data_np * (dense_mask_np != 0) ++ np.testing.assert_allclose( ++ sparse_out_pd.to_dense().numpy(), dense_data_np_ref ++ ) ++ ++ if check_grad: ++ # with sparse_out_pd backward, we get the grad from dense_data_pd ++ sparse_out_pd.backward() ++ dense_data_grad = dense_data_pd.grad ++ ++ self.assertEqual( ++ list(dense_data_grad.shape), list(dense_data_pd.shape) ++ ) ++ self.assertEqual(dense_data_grad.dtype, dense_data_pd.dtype) ++ ++ # make a dense data to compare the grad from sparse_out_pd ++ grad_ref = np.ones_like(dense_mask_np) * (dense_mask_np != 0) ++ ++ np.testing.assert_allclose( ++ dense_data_pd.grad.numpy(), ++ grad_ref, ++ ) ++ ++ def check_with_dtypes(self, shape): ++ for place in self.places: ++ self.check(shape, 'float32', place) ++ self.check(shape, 'float64', place) ++ self.check(shape, 'int32', place) ++ self.check(shape, 'int64', place) ++ self.check(shape, 'complex64', place) ++ self.check(shape, 'complex128', place) ++ ++ # `int8`` not registered in `FullLikeCooKernel`, so skip check_grad ++ self.check(shape, 'int8', place, check_grad=False) ++ ++ # `int16` not registered in `multiply`, so skip check_grad ++ self.check(shape, 'int16', place, check_grad=False) ++ ++ if paddle.is_compiled_with_cuda(): ++ place = paddle.CUDAPlace(0) ++ self.check(shape, 'float16', place) ++ ++ ++class TestMaskAsCoo(TestMaskAs): ++ def init_format(self): ++ self.format = 'coo' ++ ++ def test_1d(self): ++ self.check_with_dtypes((5,)) ++ ++ def test_2d(self): ++ self.check_with_dtypes((5, 3)) ++ ++ def test_3d(self): ++ self.check_with_dtypes((5, 3, 4)) ++ ++ def test_4d(self): ++ self.check_with_dtypes((5, 3, 4, 2)) ++ ++ ++class TestMaskAsCsr(TestMaskAs): ++ def init_format(self): ++ self.format = 'csr' ++ ++ def test_2d(self): ++ self.check_with_dtypes((5, 3)) ++ ++ def test_3d(self): ++ self.check_with_dtypes((5, 3, 4)) ++ ++ def test_error_dimension(self): ++ # error 1d ++ with self.assertRaises(ValueError): ++ self.check_with_dtypes((5,)) ++ ++ # error 4d ++ with self.assertRaises(ValueError): ++ self.check_with_dtypes((5, 3, 4, 2)) ++ ++ ++if __name__ == "__main__": diff --git a/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.sh b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.sh new file mode 100644 index 000000000..b40127d4e --- /dev/null +++ b/swe-paddle/tasks/PaddlePaddle__Paddle-64320/tests/test.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail + +python -m pytest test/legacy_test/test_sparse_mask_as_op.py -q