草稿
[WIP]concat_2d_jaggeed 前反向算子迁移 #236
tanfeng_cc创建于 5 天前
[WIP]concat_2d_jaggeed 前反向算子迁移 #236
草稿
共 29 个文件变更+4130-0
| @@ -0,0 +1,27 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# concat_2d_jagged 算子聚合 CMakeLists | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 设置 vendor 名，使用 npu_op_package + INSTALL_PATH 将产物 | ||
| 5 | +# 定向到 opp/<variant>/vendors/<op_name>/ 目录。 | ||
| 6 | +# ============================================================================ | ||
| 7 | + | ||
| 8 | +# vendor 名，同时也是 CANN package 名 | ||
| 9 | +set(VENDOR_NAME "concat_2d_jagged") | ||
| 10 | + | ||
| 11 | +# 独立 autogen 路径，避免多个算子共享冲突 | ||
| 12 | +set(ASCEND_AUTOGEN_PATH "${CMAKE_CURRENT_BINARY_DIR}/autogen") | ||
| 13 | +file(MAKE_DIRECTORY "${ASCEND_AUTOGEN_PATH}") | ||
| 14 | + | ||
| 15 | +# 安装路径 | ||
| 16 | +set(OPP_VENDOR_INSTALL_PREFIX | ||
| 17 | + "${CUSTOM_OPP_INSTALL_PREFIX}/vendors/${VENDOR_NAME}") | ||
| 18 | + | ||
| 19 | +npu_op_package(${VENDOR_NAME} | ||
| 20 | + TYPE RUN | ||
| 21 | + CONFIG | ||
| 22 | + INSTALL_PATH ${OPP_VENDOR_INSTALL_PREFIX} | ||
| 23 | + VENDOR_NAME ${VENDOR_NAME} | ||
| 24 | +) | ||
| 25 | + | ||
| 26 | +add_subdirectory(op_host) | ||
| 27 | +add_subdirectory(op_kernel) | ||
| @@ -0,0 +1,109 @@ | |||
| 1 | +# concat_2d_jagged | ||
| 2 | + | ||
| 3 | +本算子仅支持 NPU 调用。 | ||
| 4 | + | ||
| 5 | +## 产品支持情况 | ||
| 6 | + | ||
| 7 | +| 硬件型号 | 是否支持 | | ||
| 8 | +| -------------------- |------| | ||
| 9 | +| Atlas A2训练系列产品 | 是 | | ||
| 10 | +| Atlas A3训练系列产品 | 是 | | ||
| 11 | +| Atlas 推理系列产品 | 否 | | ||
| 12 | + | ||
| 13 | +## 目录结构 | ||
| 14 | + | ||
| 15 | +`concat_2d_jagged`（前向）与 `split_2d_jagged`（反向）同属于 `concat_jagged_tensor` 聚合目录，torch 插件由两个算子共用： | ||
| 16 | + | ||
| 17 | +```shell | ||
| 18 | +concat_jagged_tensor | ||
| 19 | +├── torch_extension # torch 插件（schema 注册、参数校验、autograd 注册，两算子共用） | ||
| 20 | +│ ├── concat_jagged_tensor_plugin.cpp | ||
| 21 | +│ └── concat_jagged_tensor_autograd.py | ||
| 22 | +├── concat_2d_jagged # 前向算子（本目录） | ||
| 23 | +│ ├── op_host # 算子 host 侧实现（OpDef、tiling） | ||
| 24 | +│ ├── op_kernel # 算子 kernel 侧实现 | ||
| 25 | +│ ├── examples # 调用示例（aclnn / torch） | ||
| 26 | +│ ├── test # pytest 功能测试 | ||
| 27 | +│ ├── docs # 接口文档 | ||
| 28 | +│ └── README.md # 算子说明文档（本文档） | ||
| 29 | +└── split_2d_jagged # 反向算子，按 offset 将一个 tensor 切分成两个 | ||
| 30 | + └── （目录结构同上） | ||
| 31 | +``` | ||
| 32 | + | ||
| 33 | +## 功能 | ||
| 34 | + | ||
| 35 | +将两个 jagged Tensor 按照 offset 在 dim0 维度上进行拼接，合并成一个 tensor。 | ||
| 36 | + | ||
| 37 | +对应接口： | ||
| 38 | + | ||
| 39 | +- aclnn 接口：`aclnnConcatJaggedTensor` | ||
| 40 | +- torch 接口：`ops_rec.ascendc.sequence.concat_2d_jagged` | ||
| 41 | + | ||
| 42 | +> CANN 算子名为 `ConcatJaggedTensor`，torch 算子名为 `ops_rec::concat_2d_jagged`；反向算子为 [split_2d_jagged](../split_2d_jagged/README.md)（CANN 算子名 `SplitJaggedTensor`）。 | ||
| 43 | + | ||
| 44 | +## 算子实现原理 | ||
| 45 | + | ||
| 46 | +输入: | ||
| 47 | + | ||
| 48 | +```python | ||
| 49 | +# 输入的两个tensor | ||
| 50 | +values = [tensor([[1,1,1,1,1,1,1,1], | ||
| 51 | + [2,2,2,2,2,2,2,2], | ||
| 52 | + [3,3,3,3,3,3,3,3] | ||
| 53 | + ]), | ||
| 54 | + tensor([[4,4,4,4,4,4,4,4], | ||
| 55 | + [5,5,5,5,5,5,5,5], | ||
| 56 | + [6,6,6,6,6,6,6,6], | ||
| 57 | + [7,7,7,7,7,7,7,7], | ||
| 58 | + ])] | ||
| 59 | +# 长度为2N,前N个数表示第一个tensor的offset,后N个数表示第2个tensor的offset。 | ||
| 60 | +offsets = [0, 1, 2, 3, 0, 1, 3, 4] | ||
| 61 | +# offset的长度，两个tensor的offset长度必须一致。 | ||
| 62 | +offsetLen = 4 | ||
| 63 | +# 拼接的tensor个数，这里仅支持两个。 | ||
| 64 | +jtNum = 2 | ||
| 65 | +``` | ||
| 66 | + | ||
| 67 | +输出： | ||
| 68 | + | ||
| 69 | +```python | ||
| 70 | +result = tensor([[1,1,1,1,1,1,1,1], | ||
| 71 | + [4,4,4,4,4,4,4,4], | ||
| 72 | + [2,2,2,2,2,2,2,2], | ||
| 73 | + [5,5,5,5,5,5,5,5], | ||
| 74 | + [6,6,6,6,6,6,6,6], | ||
| 75 | + [3,3,3,3,3,3,3,3], | ||
| 76 | + [7,7,7,7,7,7,7,7] | ||
| 77 | + ]) | ||
| 78 | +``` | ||
| 79 | + | ||
| 80 | +`nPrefixFromRight > 0` 时，右侧（第 2 个）tensor 每个分片的前 `nPrefixFromRight` 行会被前移到对应左侧分片之前（分片长度不足时按实际长度截断）。 | ||
| 81 | + | ||
| 82 | +算子采用 AscendC 开发。host 侧 tiling 把一次拼接拆解为 `jtNum * (offsetLen - 1)` 个连续行切片（`nPrefixFromRight > 0` 时加 1），均摊到全部 AI Vector 核（前 formerCore 个核各多承担 1 个切片）；当切片数少于核数且单个 tensor 行数超过 24K 时，再按核数将长切片等分拆细，均衡各核负载。kernel 侧每个核将自己负责的切片经 UB 中转（`DataCopyPad` 按 ubMaxLength 分段搬运），从输入 GM 搬写到输出 GM 的对应位置。 | ||
| 83 | + | ||
| 84 | +## 算子输入与输出 | ||
| 85 | + | ||
| 86 | +| 名称 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | 范围 | 说明 | | ||
| 87 | +|------------------|--------|--------|--------------------------------|------------------------|------------------------------------------------------|-----------------| | ||
| 88 | +| values | 输入 | List | bfloat16/float16/float32/int32 | List[Tensor, Tensor] | tensor_a与tensor_b必须为二维且第二维相等 | 待拼接的tensor list | | ||
| 89 | +| offsets | 输入 | List | int64 | List[int] | [4, 2048]，且列表长度需为offsetLen值的2倍，前offsetLen个为第1个tensor的offset，后offsetLen个为第2个tensor的offset | 待拼接的tensor的偏移 | | ||
| 90 | +| offsetLen | 输入 | int | int64 | NA | [2, 1024] | 单个tensor的offset长度 | | ||
| 91 | +| jtNum | 输入 | int | int64 | NA | 支持jtNum = 2 | 待拼接tensor个数 | | ||
| 92 | +| nPrefixFromRight | 输入 | int | int64 | NA | nPrefixFromRight >= 0 | 拼接时右侧tensor每个分片前移的行数 | | ||
| 93 | +| result | 输出 | Tensor | bfloat16/float16/float32/int32 | [dim0, dim1] | 结果为二维,<br/>dim0的长度等于两个tensor的dim0之和，<br/>dim1的长度与拼接tensor的dim1相等。 | NA | | ||
| 94 | + | ||
| 95 | +## 约束说明 | ||
| 96 | + | ||
| 97 | +- values 中两个 tensor 必须为 2D、dim1 相等且数据类型一致，数据类型仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 98 | +- 每个 tensor 的 offset 需单调不减、首元素为 0，末元素不超过对应 tensor 的 dim0 | ||
| 99 | +- offsetLen 取值范围 [2, 1024]，jtNum 仅支持 2，nPrefixFromRight ≥ 0 | ||
| 100 | +- 该算子需要在 NPU 设备上运行 | ||
| 101 | + | ||
| 102 | +> torch 接口的入参形式（offset_a / offset_b 以 tensor 传入）、autograd 与 torch.compile 支持及扩展层额外校验，详见 [torch 接口文档](./docs/torch_concat_2d_jagged.md)。 | ||
| 103 | + | ||
| 104 | +## 调用示例 | ||
| 105 | + | ||
| 106 | +- [aclnn 调用示例](./examples/aclnn/test_aclnn_concat_2d_jagged.cpp) | ||
| 107 | +- [torch 调用示例](./examples/torch/test_torch_concat_2d_jagged.py) | ||
| 108 | +- [aclnn 接口文档](./docs/aclnn_concat_jagged_tensor.md) | ||
| 109 | +- [torch 接口文档](./docs/torch_concat_2d_jagged.md) | ||
| @@ -0,0 +1,92 @@ | |||
| 1 | +# aclnnConcatJaggedTensor | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 硬件型号 | 是否支持 | | ||
| 6 | +| -------------------- |------| | ||
| 7 | +| Atlas A2训练系列产品 | 是 | | ||
| 8 | +| Atlas A3训练系列产品 | 是 | | ||
| 9 | +| Atlas 推理系列产品 | 否 | | ||
| 10 | + | ||
| 11 | +## 功能说明 | ||
| 12 | + | ||
| 13 | +将两个 jagged Tensor 按照 offset 在 dim0 维度上进行拼接，合并成一个 tensor。对应 torch 接口 `ops_rec.ascendc.sequence.concat_2d_jagged`，CANN 算子名为 `ConcatJaggedTensor`。 | ||
| 14 | + | ||
| 15 | +输入以动态 tensor 列表（`aclTensorList`）形式传入 2 个待拼接的 tensor；offsets 以 host 侧 int64 数组属性传入，长度为 `offsetLen` 的 2 倍，前 `offsetLen` 个为第 1 个 tensor 的 offset，后 `offsetLen` 个为第 2 个 tensor 的 offset。 | ||
| 16 | + | ||
| 17 | +设 `N = offsetLen`，第 i 个分片中 A 的行区间为 `[offsets[i], offsets[i+1])`、B 的行区间为 `[offsets[N+i], offsets[N+i+1])`： | ||
| 18 | + | ||
| 19 | +- `nPrefixFromRight = 0` 时，输出按分片交错拼接：`result = [A_0, B_0, A_1, B_1, ..., A_{N-2}, B_{N-2}]` | ||
| 20 | +- `nPrefixFromRight > 0` 时，B 每个分片的前 `nPrefixFromRight` 行（不足按实际长度截断）前移到对应 A 分片之前：`result = [B_0前缀, A_0, B_0剩余+B_1前缀, A_1, ..., B_{N-2}剩余]` | ||
| 21 | + | ||
| 22 | +## 函数原型 | ||
| 23 | + | ||
| 24 | +两段式接口，必须先调用 `aclnnConcatJaggedTensorGetWorkspaceSize` 接口获取 workspace 大小和执行器，再调用 `aclnnConcatJaggedTensor` 接口执行计算。 | ||
| 25 | + | ||
| 26 | +```c | ||
| 27 | +aclnnStatus aclnnConcatJaggedTensorGetWorkspaceSize( | ||
| 28 | + const aclTensorList *values, | ||
| 29 | + const aclIntArray *offsets, | ||
| 30 | + int64_t offsetLen, | ||
| 31 | + int64_t jtNum, | ||
| 32 | + int64_t nPrefixFromRight, | ||
| 33 | + aclTensor *out, | ||
| 34 | + uint64_t *workspaceSize, | ||
| 35 | + aclOpExecutor **executor); | ||
| 36 | + | ||
| 37 | +aclnnStatus aclnnConcatJaggedTensor( | ||
| 38 | + void *workspace, | ||
| 39 | + uint64_t workspaceSize, | ||
| 40 | + aclOpExecutor *executor, | ||
| 41 | + aclrtStream stream); | ||
| 42 | +``` | ||
| 43 | + | ||
| 44 | +## aclnnConcatJaggedTensorGetWorkspaceSize | ||
| 45 | + | ||
| 46 | +### 参数说明 | ||
| 47 | + | ||
| 48 | +| 参数名 | 输入/输出 | 描述 | 数据类型 | 数据格式 | | ||
| 49 | +|--------|---------|------|---------|---------| | ||
| 50 | +| values（aclTensorList） | 输入 | 待拼接的 2 个 tensor 组成的列表，每个 tensor shape [rows, dim1]，两个 tensor 的 dim1 必须相等 | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | | ||
| 51 | +| offsets（aclIntArray） | 输入 | host 侧偏移数组，长度为 `2 * offsetLen`，前 offsetLen 个为第 1 个 tensor 的 offset，后 offsetLen 个为第 2 个 tensor 的 offset | INT64 | - | | ||
| 52 | +| offsetLen（int64_t） | 输入 | 单个 tensor 的 offset 长度，取值范围 [2, 1024] | - | - | | ||
| 53 | +| jtNum（int64_t） | 输入 | 待拼接 tensor 个数，仅支持 2 | - | - | | ||
| 54 | +| nPrefixFromRight（int64_t） | 输入 | 拼接时右侧 tensor 每个分片前移的行数，取值 ≥ 0 | - | - | | ||
| 55 | +| out（aclTensor） | 输出 | 拼接结果，shape [rowsA + rowsB, dim1]，数据类型与输入一致 | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | | ||
| 56 | +| workspaceSize（uint64_t*） | 输出 | 返回需要在 Device 侧申请的 workspace 大小 | - | - | | ||
| 57 | +| executor（aclOpExecutor**） | 输出 | 返回算子执行器 | - | - | | ||
| 58 | + | ||
| 59 | +### 返回值 | ||
| 60 | + | ||
| 61 | +| 返回值 | 错误码 | 描述 | | ||
| 62 | +|--------|-------|------| | ||
| 63 | +| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入参数是必选输入/输出且为空指针 | | ||
| 64 | +| ACLNN_ERR_PARAM_INVALID | 161002 | 参数数据类型或数据格式不在支持范围内 | | ||
| 65 | + | ||
| 66 | +## aclnnConcatJaggedTensor | ||
| 67 | + | ||
| 68 | +### 参数说明 | ||
| 69 | + | ||
| 70 | +| 参数名 | 输入/输出 | 描述 | | ||
| 71 | +|--------|---------|------| | ||
| 72 | +| workspace | 输入 | 在 Device 侧申请的 workspace 内存地址 | | ||
| 73 | +| workspaceSize | 输入 | workspace 大小，由 `aclnnConcatJaggedTensorGetWorkspaceSize` 获取 | | ||
| 74 | +| executor | 输入 | 算子执行器 | | ||
| 75 | +| stream | 输入 | 指定执行任务的 Stream | | ||
| 76 | + | ||
| 77 | +### 返回值 | ||
| 78 | + | ||
| 79 | +aclnnStatus：返回状态码，成功返回 `ACLNN_SUCCESS`。 | ||
| 80 | + | ||
| 81 | +## 约束说明 | ||
| 82 | + | ||
| 83 | +- values 列表长度必须等于 jtNum（仅支持 2），两个 tensor 必须为 2D、dim1 相等且数据类型一致 | ||
| 84 | +- 输入数据类型仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 85 | +- 每个 tensor 的 offset 需单调不减，首元素为 0，末元素不超过对应 tensor 的 dim0 | ||
| 86 | +- offsetLen 取值范围 [2, 1024]，nPrefixFromRight 取值 ≥ 0 | ||
| 87 | +- out 的 shape 必须为 [values[0].dim0 + values[1].dim0, values[0].dim1]，需由调用侧提前分配 | ||
| 88 | +- 该算子需要在 NPU 设备上运行 | ||
| 89 | + | ||
| 90 | +## 调用示例 | ||
| 91 | + | ||
| 92 | +具体参考 [test_aclnn_concat_2d_jagged.cpp](../examples/aclnn/test_aclnn_concat_2d_jagged.cpp) 的完整实现。 | ||
| @@ -0,0 +1,85 @@ | |||
| 1 | +# concat_2d_jagged | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 硬件型号 | 是否支持 | | ||
| 6 | +| -------------------- |------| | ||
| 7 | +| Atlas A2训练系列产品 | 是 | | ||
| 8 | +| Atlas A3训练系列产品 | 是 | | ||
| 9 | +| Atlas 推理系列产品 | 否 | | ||
| 10 | + | ||
| 11 | +## 功能说明 | ||
| 12 | + | ||
| 13 | +将两个 jagged Tensor 按照 offset 在 dim0 维度上进行拼接，合并成一个 tensor。底层调用 aclnn 接口 `aclnnConcatJaggedTensor`（CANN 算子名 `ConcatJaggedTensor`）。 | ||
| 14 | + | ||
| 15 | +设 `N = offset_a.shape[0]`，第 i 个分片中 A 的行区间为 `[offset_a[i], offset_a[i+1])`、B 的行区间为 `[offset_b[i], offset_b[i+1])`： | ||
| 16 | + | ||
| 17 | +- `n_prefix_from_right = 0` 时，输出按分片交错拼接：`result = [A_0, B_0, A_1, B_1, ..., A_{N-2}, B_{N-2}]` | ||
| 18 | +- `n_prefix_from_right > 0` 时，B 每个分片的前 `n_prefix_from_right` 行（不足按实际长度截断）前移到对应 A 分片之前 | ||
| 19 | + | ||
| 20 | +## 接口原型 | ||
| 21 | + | ||
| 22 | +```python | ||
| 23 | +ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 24 | + max_seqlen: int, | ||
| 25 | + values_a: Tensor, | ||
| 26 | + values_b: Tensor, | ||
| 27 | + offset_a: Tensor, | ||
| 28 | + offset_b: Tensor, | ||
| 29 | + is_replace: bool = False, | ||
| 30 | + n_prefix_from_right: int = 0, | ||
| 31 | +) -> Tensor | ||
| 32 | +``` | ||
| 33 | + | ||
| 34 | +## 参数说明 | ||
| 35 | + | ||
| 36 | +| 参数名 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | shape | 说明 | | ||
| 37 | +|--------|---------|---------|---------|---------|-------|------| | ||
| 38 | +| max_seqlen | 输入 | int | - | - | - | 保留参数，当前实现未使用 | | ||
| 39 | +| values_a | 输入 | Tensor | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | [rows_a, dim1] | 待拼接的第 1 个 tensor | | ||
| 40 | +| values_b | 输入 | Tensor | 与 values_a 一致 | ND | [rows_b, dim1] | 待拼接的第 2 个 tensor，dim1 必须与 values_a 相等 | | ||
| 41 | +| offset_a | 输入 | Tensor | INT64 / INT32 | ND | [N] | values_a 的 offset，单调不减、首元素为 0 | | ||
| 42 | +| offset_b | 输入 | Tensor | INT64 / INT32 | ND | [N] | values_b 的 offset，长度必须与 offset_a 一致 | | ||
| 43 | +| is_replace | 输入 | bool | - | - | - | 保留参数，当前实现未使用 | | ||
| 44 | +| n_prefix_from_right | 输入 | int | - | - | - | 拼接时 values_b 每个分片前移的行数，取值 ≥ 0 | | ||
| 45 | +| output | 输出 | Tensor | 与 values_a 一致 | ND | [rows_a + rows_b, dim1] | 拼接结果 | | ||
| 46 | + | ||
| 47 | +> offset_a / offset_b 以 NPU 上的 tensor 形式传入，插件内部会将其拷贝到 host 并拼接成长度为 2N 的 int64 数组传给底层算子。 | ||
| 48 | + | ||
| 49 | +## autograd 与 torch.compile | ||
| 50 | + | ||
| 51 | +- **自动求导**：concat 的反向等价于按相同 offset 对 `grad_output` 做 split。由 [concat_jagged_tensor_autograd.py](../../torch_extension/concat_jagged_tensor_autograd.py) 注册，反向内部调用 `ops_rec.ascendc.sequence.split_2d_jagged`，对 values_a / values_b 返回切分后的梯度，其余入参（offset 及标量参数）不需要梯度。 | ||
| 52 | +- **torch.compile**：注册了 Meta 实现，输出 shape 由两个输入的 dim0 之和静态决定（`empty_symint` 实现，不强制特化动态 shape），支持 dynamo / inductor 入图（`dynamic=True, fullgraph=True`）。 | ||
| 53 | + | ||
| 54 | +## 约束说明 | ||
| 55 | + | ||
| 56 | +约束分为两层：算子自身语义适用于 aclnn 与 torch 两条调用路径；torch 扩展层额外加的入参校验仅在调用 `ops_rec.ascendc.sequence.concat_2d_jagged` 时生效。 | ||
| 57 | + | ||
| 58 | +### 算子约束 | ||
| 59 | + | ||
| 60 | +- values_a / values_b 必须为 2D、dim1 相等且数据类型一致，仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 61 | +- offset_a / offset_b 必须为 1D、长度相等，取值范围 [2, 1024]，dtype 为 INT64 / INT32 | ||
| 62 | +- 每个 tensor 的 offset 需单调不减，末元素不超过对应 tensor 的 dim0 | ||
| 63 | +- 一次仅支持拼接 2 个 tensor（jtNum = 2） | ||
| 64 | +- 该算子需要在 NPU 设备上运行 | ||
| 65 | + | ||
| 66 | +### torch 扩展层额外校验 | ||
| 67 | + | ||
| 68 | +由 [concat_jagged_tensor_plugin.cpp](../../torch_extension/concat_jagged_tensor_plugin.cpp) 在调用算子前完成： | ||
| 69 | + | ||
| 70 | +- 4 个输入 tensor 必须已定义且位于同一 NPU 设备 | ||
| 71 | +- values_a / values_b 必须为 2D 且 dim1 相等，dtype 仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 且两者一致 | ||
| 72 | +- offset_a / offset_b 必须为 1D，dtype 为 INT64 / INT32，长度相等且在 [2, 1024] 范围内 | ||
| 73 | +- `values_a.shape[0] >= offset_a[-1]`，`values_b.shape[0] >= offset_b[-1]`，且 `values_a.shape[0] >= 1` | ||
| 74 | +- `n_prefix_from_right >= 0` | ||
| 75 | +- 需要安装 `torch_npu` | ||
| 76 | + | ||
| 77 | +## 安装方式 | ||
| 78 | + | ||
| 79 | +```bash | ||
| 80 | +bash build.sh --install -j16 | ||
| 81 | +``` | ||
| 82 | + | ||
| 83 | +## 调用示例 | ||
| 84 | + | ||
| 85 | +完整调用示例（含前向、nPrefixFromRight、autograd、torch.compile 用例）参考 [test_torch_concat_2d_jagged.py](../examples/torch/test_torch_concat_2d_jagged.py)。 | ||
| @@ -0,0 +1,254 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | +limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +/** | ||
| 17 | + * @file test_aclnn_concat_2d_jagged.cpp | ||
| 18 | + * @brief 通过 aclnnConcatJaggedTensor 接口直调 concat_2d_jagged 算子 | ||
| 19 | + */ | ||
| 20 | + | ||
| 21 | +#include <cmath> | ||
| 22 | +#include <cstdint> | ||
| 23 | +#include <cstdio> | ||
| 24 | +#include <vector> | ||
| 25 | +#include "acl/acl.h" | ||
| 26 | +#include "aclnn_concat_jagged_tensor.h" | ||
| 27 | + | ||
| 28 | +#define CHECK_RET(cond, return_expr) \ | ||
| 29 | + do { \ | ||
| 30 | + if (!(cond)) { \ | ||
| 31 | + return_expr; \ | ||
| 32 | + } \ | ||
| 33 | + } while (0) | ||
| 34 | + | ||
| 35 | +#define LOG_PRINT(message, ...) \ | ||
| 36 | + do { \ | ||
| 37 | + printf(message, ##__VA_ARGS__); \ | ||
| 38 | + } while (0) | ||
| 39 | + | ||
| 40 | +constexpr int64_t JT_NUM = 2; // 一次拼接 2 个 tensor | ||
| 41 | + | ||
| 42 | +int64_t GetShapeSize(const std::vector<int64_t>& shape) | ||
| 43 | +{ | ||
| 44 | + int64_t shape_size = 1; | ||
| 45 | + for (auto i : shape) { | ||
| 46 | + shape_size *= i; | ||
| 47 | + } | ||
| 48 | + return shape_size; | ||
| 49 | +} | ||
| 50 | + | ||
| 51 | +int Init(int32_t deviceId, aclrtStream* stream) | ||
| 52 | +{ | ||
| 53 | + auto ret = aclInit(nullptr); | ||
| 54 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); | ||
| 55 | + ret = aclrtSetDevice(deviceId); | ||
| 56 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); | ||
| 57 | + ret = aclrtCreateStream(stream); | ||
| 58 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); | ||
| 59 | + return 0; | ||
| 60 | +} | ||
| 61 | + | ||
| 62 | +template <typename T> | ||
| 63 | +int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, | ||
| 64 | + aclDataType dataType, aclTensor** tensor) | ||
| 65 | +{ | ||
| 66 | + auto size = GetShapeSize(shape) * sizeof(T); | ||
| 67 | + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 68 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); | ||
| 69 | + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 70 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); | ||
| 71 | + | ||
| 72 | + std::vector<int64_t> strides(shape.size(), 1); | ||
| 73 | + for (int64_t i = static_cast<int64_t>(shape.size()) - 2; i >= 0; --i) { | ||
| 74 | + strides[i] = shape[i + 1] * strides[i + 1]; | ||
| 75 | + } | ||
| 76 | + | ||
| 77 | + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, | ||
| 78 | + shape.data(), shape.size(), *deviceAddr); | ||
| 79 | + return 0; | ||
| 80 | +} | ||
| 81 | + | ||
| 82 | +// 执行一次 concat 前向：GetWorkspaceSize -> 执行 -> 同步（workspace 在函数内分配与释放） | ||
| 83 | +int RunConcat(aclTensorList* valuesList, aclIntArray* offsetsArray, int64_t offsetLen, | ||
| 84 | + int64_t nPrefixFromRight, aclTensor* out, aclrtStream stream) | ||
| 85 | +{ | ||
| 86 | + uint64_t workspaceSize = 0; | ||
| 87 | + aclOpExecutor* executor = nullptr; | ||
| 88 | + auto ret = aclnnConcatJaggedTensorGetWorkspaceSize(valuesList, offsetsArray, offsetLen, JT_NUM, | ||
| 89 | + nPrefixFromRight, out, &workspaceSize, &executor); | ||
| 90 | + CHECK_RET(ret == ACL_SUCCESS, | ||
| 91 | + LOG_PRINT("aclnnConcatJaggedTensorGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); | ||
| 92 | + | ||
| 93 | + void* workspaceAddr = nullptr; | ||
| 94 | + if (workspaceSize != 0) { | ||
| 95 | + ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 96 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc workspace failed. ERROR: %d\n", ret); return ret); | ||
| 97 | + } | ||
| 98 | + ret = aclnnConcatJaggedTensor(workspaceAddr, workspaceSize, executor, stream); | ||
| 99 | + if (workspaceAddr != nullptr) { | ||
| 100 | + aclrtFree(workspaceAddr); | ||
| 101 | + } | ||
| 102 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnConcatJaggedTensor failed. ERROR: %d\n", ret); return ret); | ||
| 103 | + | ||
| 104 | + ret = aclrtSynchronizeDevice(); | ||
| 105 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeDevice failed. ERROR: %d\n", ret); return ret); | ||
| 106 | + return 0; | ||
| 107 | +} | ||
| 108 | + | ||
| 109 | +// 逐元素比较，返回错误元素个数 | ||
| 110 | +int64_t CompareBuf(const std::vector<float>& actual, const std::vector<float>& expected, float tol = 1e-6f) | ||
| 111 | +{ | ||
| 112 | + if (actual.size() != expected.size()) { | ||
| 113 | + LOG_PRINT("size mismatch: actual=%zu expected=%zu\n", actual.size(), expected.size()); | ||
| 114 | + return -1; | ||
| 115 | + } | ||
| 116 | + int64_t errCnt = 0; | ||
| 117 | + for (size_t i = 0; i < actual.size(); ++i) { | ||
| 118 | + if (std::fabs(actual[i] - expected[i]) > tol) { | ||
| 119 | + ++errCnt; | ||
| 120 | + } | ||
| 121 | + } | ||
| 122 | + return errCnt; | ||
| 123 | +} | ||
| 124 | + | ||
| 125 | +// 把 {v1, v2, ...} 各扩展为 cols 个元素的一行 | ||
| 126 | +std::vector<float> RepeatRows(const std::vector<float>& rowValues, int64_t cols) | ||
| 127 | +{ | ||
| 128 | + std::vector<float> out; | ||
| 129 | + for (float v : rowValues) { | ||
| 130 | + out.insert(out.end(), cols, v); | ||
| 131 | + } | ||
| 132 | + return out; | ||
| 133 | +} | ||
| 134 | + | ||
| 135 | +int main() | ||
| 136 | +{ | ||
| 137 | + int32_t deviceId = 0; | ||
| 138 | + aclrtStream stream; | ||
| 139 | + auto ret = Init(deviceId, &stream); | ||
| 140 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); | ||
| 141 | + | ||
| 142 | + int failures = 0; | ||
| 143 | + | ||
| 144 | + // ============================================================== | ||
| 145 | + // Case 1 & 2（对齐 torch: test_concat_basic / test_concat_with_prefix） | ||
| 146 | + // VALUES_A (3, 8): 每行全 1 / 2 / 3 | ||
| 147 | + // VALUES_B (4, 8): 每行全 4 / 5 / 6 / 7 | ||
| 148 | + // OFFSET_A = [0, 1, 2, 3], OFFSET_B = [0, 1, 3, 4] | ||
| 149 | + // ============================================================== | ||
| 150 | + constexpr int64_t kRowsA = 3; | ||
| 151 | + constexpr int64_t kRowsB = 4; | ||
| 152 | + constexpr int64_t kCols = 8; | ||
| 153 | + constexpr int64_t kTotalRows = kRowsA + kRowsB; // 7 | ||
| 154 | + constexpr int64_t kOffsetLen = 4; | ||
| 155 | + | ||
| 156 | + std::vector<float> valuesAHost = RepeatRows({1.0f, 2.0f, 3.0f}, kCols); | ||
| 157 | + std::vector<float> valuesBHost = RepeatRows({4.0f, 5.0f, 6.0f, 7.0f}, kCols); | ||
| 158 | + std::vector<int64_t> offsetsHost = {0, 1, 2, 3, 0, 1, 3, 4}; | ||
| 159 | + | ||
| 160 | + void* valuesADeviceAddr = nullptr; | ||
| 161 | + void* valuesBDeviceAddr = nullptr; | ||
| 162 | + void* outDeviceAddr = nullptr; | ||
| 163 | + aclTensor* valuesA = nullptr; | ||
| 164 | + aclTensor* valuesB = nullptr; | ||
| 165 | + aclTensor* out = nullptr; | ||
| 166 | + | ||
| 167 | + ret = CreateAclTensor(valuesAHost, {kRowsA, kCols}, &valuesADeviceAddr, aclDataType::ACL_FLOAT, &valuesA); | ||
| 168 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 169 | + ret = CreateAclTensor(valuesBHost, {kRowsB, kCols}, &valuesBDeviceAddr, aclDataType::ACL_FLOAT, &valuesB); | ||
| 170 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 171 | + | ||
| 172 | + // 动态输入：将 A、B 组装为 aclTensorList | ||
| 173 | + aclTensor* valueTensors[JT_NUM] = {valuesA, valuesB}; | ||
| 174 | + aclTensorList* valuesList = aclCreateTensorList(valueTensors, JT_NUM); | ||
| 175 | + CHECK_RET(valuesList != nullptr, LOG_PRINT("aclCreateTensorList failed.\n"); return 1); | ||
| 176 | + aclIntArray* offsetsArray = aclCreateIntArray(offsetsHost.data(), offsetsHost.size()); | ||
| 177 | + CHECK_RET(offsetsArray != nullptr, LOG_PRINT("aclCreateIntArray failed.\n"); return 1); | ||
| 178 | + | ||
| 179 | + // -------------------------------------------------------------- | ||
| 180 | + // Case 1: nPrefixFromRight = 0（对齐 test_concat_basic） | ||
| 181 | + // -------------------------------------------------------------- | ||
| 182 | + { | ||
| 183 | + std::vector<float> outHost(kTotalRows * kCols, 0); | ||
| 184 | + ret = CreateAclTensor(outHost, {kTotalRows, kCols}, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); | ||
| 185 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 186 | + | ||
| 187 | + ret = RunConcat(valuesList, offsetsArray, kOffsetLen, /*nPrefixFromRight=*/0, out, stream); | ||
| 188 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 189 | + | ||
| 190 | + auto outSize = kTotalRows * kCols * sizeof(float); | ||
| 191 | + ret = aclrtMemcpy(outHost.data(), outSize, outDeviceAddr, outSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 192 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy output failed. ERROR: %d\n", ret); return ret); | ||
| 193 | + | ||
| 194 | + auto expected = RepeatRows({1.0f, 4.0f, 2.0f, 5.0f, 6.0f, 3.0f, 7.0f}, kCols); | ||
| 195 | + int64_t errCnt = CompareBuf(outHost, expected); | ||
| 196 | + if (errCnt != 0) { | ||
| 197 | + LOG_PRINT("[FAIL] concat basic (nPrefix=0), errCnt=%ld\n", errCnt); | ||
| 198 | + ++failures; | ||
| 199 | + } else { | ||
| 200 | + LOG_PRINT("[PASS] concat basic (nPrefix=0)\n"); | ||
| 201 | + } | ||
| 202 | + aclDestroyTensor(out); | ||
| 203 | + aclrtFree(outDeviceAddr); | ||
| 204 | + out = nullptr; | ||
| 205 | + outDeviceAddr = nullptr; | ||
| 206 | + } | ||
| 207 | + | ||
| 208 | + // -------------------------------------------------------------- | ||
| 209 | + // Case 2: nPrefixFromRight = 1（对齐 test_concat_with_prefix） | ||
| 210 | + // -------------------------------------------------------------- | ||
| 211 | + { | ||
| 212 | + std::vector<float> outHost(kTotalRows * kCols, 0); | ||
| 213 | + ret = CreateAclTensor(outHost, {kTotalRows, kCols}, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); | ||
| 214 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 215 | + | ||
| 216 | + ret = RunConcat(valuesList, offsetsArray, kOffsetLen, /*nPrefixFromRight=*/1, out, stream); | ||
| 217 | + CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 218 | + | ||
| 219 | + auto outSize = kTotalRows * kCols * sizeof(float); | ||
| 220 | + ret = aclrtMemcpy(outHost.data(), outSize, outDeviceAddr, outSize, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 221 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy output failed. ERROR: %d\n", ret); return ret); | ||
| 222 | + | ||
| 223 | + auto expected = RepeatRows({4.0f, 1.0f, 5.0f, 2.0f, 6.0f, 7.0f, 3.0f}, kCols); | ||
| 224 | + int64_t errCnt = CompareBuf(outHost, expected); | ||
| 225 | + if (errCnt != 0) { | ||
| 226 | + LOG_PRINT("[FAIL] concat with prefix (nPrefix=1), errCnt=%ld\n", errCnt); | ||
| 227 | + ++failures; | ||
| 228 | + } else { | ||
| 229 | + LOG_PRINT("[PASS] concat with prefix (nPrefix=1)\n"); | ||
| 230 | + } | ||
| 231 | + aclDestroyTensor(out); | ||
| 232 | + aclrtFree(outDeviceAddr); | ||
| 233 | + out = nullptr; | ||
| 234 | + outDeviceAddr = nullptr; | ||
| 235 | + } | ||
| 236 | + | ||
| 237 | + aclDestroyTensorList(valuesList); | ||
| 238 | + aclDestroyIntArray(offsetsArray); | ||
| 239 | + aclDestroyTensor(valuesA); | ||
| 240 | + aclDestroyTensor(valuesB); | ||
| 241 | + aclrtFree(valuesADeviceAddr); | ||
| 242 | + aclrtFree(valuesBDeviceAddr); | ||
| 243 | + | ||
| 244 | + aclrtDestroyStream(stream); | ||
| 245 | + aclrtResetDevice(deviceId); | ||
| 246 | + aclFinalize(); | ||
| 247 | + | ||
| 248 | + if (failures > 0) { | ||
| 249 | + LOG_PRINT("=== concat_2d_jagged aclnn example FAILED, %d case(s) failed ===\n", failures); | ||
| 250 | + return 1; | ||
| 251 | + } | ||
| 252 | + LOG_PRINT("[aclnn] concat_2d_jagged aclnn example PASSED\n"); | ||
| 253 | + return 0; | ||
| 254 | +} | ||
Aascendc/sequence/concat_jagged_tensor/concat_2d_jagged/examples/torch/test_torch_concat_2d_jagged.py+165-0
| @@ -0,0 +1,165 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | +# ============================================================================== | ||
| 16 | + | ||
| 17 | +""" | ||
| 18 | +concat_2d_jagged 算子 PyTorch 接口基础用例 | ||
| 19 | + | ||
| 20 | +演示: | ||
| 21 | + 1. concat_2d_jagged 基本前向拼接 (nPrefixFromRight=0) | ||
| 22 | + 2. nPrefixFromRight 参数效果 (n_prefix in [1, 2]) | ||
| 23 | + 3. concat 前向 + backward 自动求导 | ||
| 24 | + 4. torch.compile / inductor 入图基础用例 | ||
| 25 | + | ||
| 26 | +用法: | ||
| 27 | + python3 test_concat_2d_jagged.py | ||
| 28 | +""" | ||
| 29 | + | ||
| 30 | +import torch | ||
| 31 | +import ops_rec | ||
| 32 | + | ||
| 33 | +DEVICE = "npu:0" | ||
| 34 | +torch.npu.config.allow_internal_format = False | ||
| 35 | +torch.npu.set_device(DEVICE) | ||
| 36 | + | ||
| 37 | +# ============================================================================== | ||
| 38 | +# 公共数据 (被 test_concat_basic / test_concat_with_prefix 复用) | ||
| 39 | +# ============================================================================== | ||
| 40 | + | ||
| 41 | +# 每行一个值，便于直观阅读（注释里标注每个值对应的 chunk） | ||
| 42 | +VALUES_A = torch.tensor([ | ||
| 43 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 44 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 45 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 46 | +], dtype=torch.float32) | ||
| 47 | + | ||
| 48 | +VALUES_B = torch.tensor([ | ||
| 49 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] | ||
| 50 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] | ||
| 51 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] | ||
| 52 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] | ||
| 53 | +], dtype=torch.float32) | ||
| 54 | + | ||
| 55 | +OFFSET_A = torch.tensor([0, 1, 2, 3], dtype=torch.int64) | ||
| 56 | +OFFSET_B = torch.tensor([0, 1, 3, 4], dtype=torch.int64) | ||
| 57 | + | ||
| 58 | +# ============================================================================== | ||
| 59 | +# 测试用例 | ||
| 60 | +# ============================================================================== | ||
| 61 | + | ||
| 62 | +def test_concat_basic(): | ||
| 63 | + """基本前向：nPrefixFromRight=0，交错拼接小样本""" | ||
| 64 | + EXPECTED_BASIC = torch.tensor([ | ||
| 65 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 66 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] | ||
| 67 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 68 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] | ||
| 69 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] | ||
| 70 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 71 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] | ||
| 72 | + ], dtype=torch.float32) | ||
| 73 | + | ||
| 74 | + result = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 75 | + 0, VALUES_A.to(DEVICE), VALUES_B.to(DEVICE), OFFSET_A.to(DEVICE), OFFSET_B.to(DEVICE), False, 0 | ||
| 76 | + ) | ||
| 77 | + diff = (result.cpu() - EXPECTED_BASIC).abs().max().item() | ||
| 78 | + assert diff < 1e-6 | ||
| 79 | + print(f"[torch] concat basic (nPrefix=0) PASSED: max diff = {diff:.6f}") | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +def test_concat_with_prefix(): | ||
| 83 | + """验证 nPrefixFromRight = 1 的效果""" | ||
| 84 | + EXPECTED_PREFIX_1 = torch.tensor([ | ||
| 85 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] 第 1 行前移 | ||
| 86 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 87 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] 第 1 行前移 | ||
| 88 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 89 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] 剩余 | ||
| 90 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] 第 1 行前移 | ||
| 91 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 92 | + ], dtype=torch.float32) | ||
| 93 | + | ||
| 94 | + # nPrefix=1 | ||
| 95 | + result = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 96 | + 0, VALUES_A.to(DEVICE), VALUES_B.to(DEVICE), OFFSET_A.to(DEVICE), OFFSET_B.to(DEVICE), False, 1 | ||
| 97 | + ) | ||
| 98 | + diff = (result.cpu() - EXPECTED_PREFIX_1).abs().max().item() | ||
| 99 | + assert diff < 1e-6 | ||
| 100 | + print(f"[torch] concat with prefix (nPrefix=1) PASSED: max diff = {diff:.6f}") | ||
| 101 | + | ||
| 102 | + | ||
| 103 | +def test_concat_autograd(): | ||
| 104 | + """concat_2d_jagged 前向 + backward 自动求导 | ||
| 105 | + | ||
| 106 | + 验证 concat.backward 通过 autograd.py 注册的 split_2d_jagged 反向路径 | ||
| 107 | + 正确把 grad_output 切回 (grad_a, grad_b)。 | ||
| 108 | + """ | ||
| 109 | + torch.manual_seed(42) | ||
| 110 | + | ||
| 111 | + rows_a, rows_b, cols = 10, 15, 8 | ||
| 112 | + values_a = torch.randn(rows_a, cols, dtype=torch.float32, device=DEVICE, requires_grad=True) | ||
| 113 | + values_b = torch.randn(rows_b, cols, dtype=torch.float32, device=DEVICE, requires_grad=True) | ||
| 114 | + offset_a = torch.tensor([0, 3, 7, 10], dtype=torch.int64).npu() | ||
| 115 | + offset_b = torch.tensor([0, 4, 9, 15], dtype=torch.int64).npu() | ||
| 116 | + | ||
| 117 | + concat_out = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 118 | + 0, values_a.to(DEVICE), values_b.to(DEVICE), offset_a.to(DEVICE), offset_b.to(DEVICE), False, 0 | ||
| 119 | + ) | ||
| 120 | + concat_out.sum().backward() | ||
| 121 | + | ||
| 122 | + assert torch.allclose(values_a.grad, torch.ones_like(values_a)), \ | ||
| 123 | + f"values_a_grad max diff = {(values_a.grad - 1).abs().max().item()}" | ||
| 124 | + assert torch.allclose(values_b.grad, torch.ones_like(values_b.grad)), \ | ||
| 125 | + f"values_b_grad max diff = {(values_b.grad - 1).abs().max().item()}" | ||
| 126 | + print("[torch] concat autograd PASSED") | ||
| 127 | + | ||
| 128 | + | ||
| 129 | +def test_concat_inductor(): | ||
| 130 | + """torch.compile / inductor 入图基础用例 | ||
| 131 | + | ||
| 132 | + 验证算子可被 dynamo 跟踪进图（dynamic=True, fullgraph=True）， | ||
| 133 | + 且 compiled 输出与 eager 一致。 | ||
| 134 | + """ | ||
| 135 | + torch.manual_seed(42) | ||
| 136 | + | ||
| 137 | + rows_a, rows_b, cols = 50, 80, 128 | ||
| 138 | + values_a = torch.randn(rows_a, cols, dtype=torch.float32, device=DEVICE) | ||
| 139 | + values_b = torch.randn(rows_b, cols, dtype=torch.float32, device=DEVICE) | ||
| 140 | + offset_a = torch.tensor([0, 15, 30, 50], dtype=torch.int64).npu() | ||
| 141 | + offset_b = torch.tensor([0, 20, 45, 80], dtype=torch.int64).npu() | ||
| 142 | + | ||
| 143 | + eager_result = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 144 | + 0, values_a.to(DEVICE), values_b.to(DEVICE), offset_a.to(DEVICE), offset_b.to(DEVICE), False, 0 | ||
| 145 | + ) | ||
| 146 | + | ||
| 147 | + compiled_fn = torch.compile( | ||
| 148 | + ops_rec.ascendc.sequence.concat_2d_jagged, | ||
| 149 | + dynamic=True, fullgraph=True, | ||
| 150 | + ) | ||
| 151 | + compiled_result = compiled_fn(0, values_a.to(DEVICE), values_b.to(DEVICE), offset_a.to(DEVICE), offset_b.to(DEVICE), False, 0) | ||
| 152 | + | ||
| 153 | + diff = (compiled_result.cpu() - eager_result.cpu()).abs().max().item() | ||
| 154 | + print(f"[torch] concat inductor: eager.shape={eager_result.shape}, compiled.shape={compiled_result.shape}") | ||
| 155 | + print(f"[torch] concat inductor: max diff = {diff:.6f}") | ||
| 156 | + assert diff < 1e-4 | ||
| 157 | + print("[torch] concat inductor PASSED") | ||
| 158 | + | ||
| 159 | + | ||
| 160 | +if __name__ == "__main__": | ||
| 161 | + test_concat_basic() | ||
| 162 | + test_concat_with_prefix() | ||
| 163 | + test_concat_autograd() | ||
| 164 | + test_concat_inductor() | ||
| 165 | + print("\n=== All concat_2d_jagged examples PASSED ===") | ||
| @@ -0,0 +1,78 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# concat_2d_jagged/op_host：算子 host 侧实现（OpDef + tiling + infershape） | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 使用 CANN 的 npu_op_code_gen + npu_op_library 编译。 | ||
| 5 | +# npu_op_package 已在上层 CMakeLists.txt 调用。 | ||
| 6 | +set(MY_COMMON_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../../../common) | ||
| 7 | + | ||
| 8 | +# OPS_REC_SOC_CXX_DEFS 由根 CMakeLists.txt 统一设置（SUPPORT_950 / SUPPORT_A2 / SUPPORT_A3） | ||
| 9 | + | ||
| 10 | +file(GLOB ops_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) | ||
| 11 | +file(GLOB ops_cc_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cc) | ||
| 12 | +list(APPEND ops_srcs ${ops_cc_srcs}) | ||
| 13 | + | ||
| 14 | +# 收集 SOC 架构特化源码（op_host/atlas950/ 或 op_host/atlasA2/ 目录下的额外实现） | ||
| 15 | +ops_rec_collect_arch_sources(ops_srcs) | ||
| 16 | + | ||
| 17 | +npu_op_code_gen( | ||
| 18 | + SRC ${ops_srcs} | ||
| 19 | + PACKAGE ${VENDOR_NAME} | ||
| 20 | + OUT_DIR ${ASCEND_AUTOGEN_PATH} | ||
| 21 | + COMPILE_OPTIONS | ||
| 22 | + ${OPS_REC_SOC_CXX_DEFS} | ||
| 23 | + -I${MY_COMMON_DIR}/log | ||
| 24 | + -I${MY_COMMON_DIR}/utils | ||
| 25 | +) | ||
| 26 | + | ||
| 27 | +file(GLOB autogen_aclnn_src ${ASCEND_AUTOGEN_PATH}/aclnn_*.cpp) | ||
| 28 | +set_source_files_properties(${autogen_aclnn_src} PROPERTIES GENERATED TRUE) | ||
| 29 | +npu_op_library(cust_opapi ACLNN | ||
| 30 | + ${autogen_aclnn_src} | ||
| 31 | +) | ||
| 32 | + | ||
| 33 | +target_include_directories(cust_opapi PRIVATE | ||
| 34 | + ${MY_COMMON_DIR}/log | ||
| 35 | + ${MY_COMMON_DIR}/utils | ||
| 36 | +) | ||
| 37 | + | ||
| 38 | +target_compile_options(cust_opapi PRIVATE | ||
| 39 | + -fvisibility=hidden | ||
| 40 | +) | ||
| 41 | + | ||
| 42 | +file(GLOB group_proto_src ${ASCEND_AUTOGEN_PATH}/group_op_proto/*.cc) | ||
| 43 | +file(GLOB proto_src ${ASCEND_AUTOGEN_PATH}/op_proto.cc) | ||
| 44 | +set_source_files_properties(${group_proto_src} PROPERTIES GENERATED TRUE) | ||
| 45 | +set_source_files_properties(${proto_src} PROPERTIES GENERATED TRUE) | ||
| 46 | +npu_op_library(cust_op_proto GRAPH | ||
| 47 | + ${ops_srcs} | ||
| 48 | + ${group_proto_src} | ||
| 49 | + ${proto_src} | ||
| 50 | +) | ||
| 51 | + | ||
| 52 | +target_include_directories(cust_op_proto PRIVATE | ||
| 53 | + ${MY_COMMON_DIR}/log | ||
| 54 | + ${MY_COMMON_DIR}/utils | ||
| 55 | +) | ||
| 56 | + | ||
| 57 | +target_compile_options(cust_op_proto PRIVATE | ||
| 58 | + -fvisibility=hidden | ||
| 59 | +) | ||
| 60 | + | ||
| 61 | +file(GLOB fallback_src ${ASCEND_AUTOGEN_PATH}/fallback_*.cpp) | ||
| 62 | +set_source_files_properties(${fallback_src} PROPERTIES GENERATED TRUE) | ||
| 63 | +npu_op_library(cust_optiling TILING | ||
| 64 | + ${ops_srcs} | ||
| 65 | + ${fallback_src} | ||
| 66 | +) | ||
| 67 | + | ||
| 68 | +target_include_directories(cust_optiling PRIVATE | ||
| 69 | + ${MY_COMMON_DIR}/log | ||
| 70 | + ${MY_COMMON_DIR}/utils | ||
| 71 | +) | ||
| 72 | + | ||
| 73 | +npu_op_package_add(${VENDOR_NAME} | ||
| 74 | + LIBRARY | ||
| 75 | + cust_optiling | ||
| 76 | + cust_opapi | ||
| 77 | + cust_op_proto | ||
| 78 | +) | ||
| @@ -0,0 +1,302 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +#include "concat_jagged_tensor_tiling.h" | ||
| 17 | +#include "register/op_def_registry.h" | ||
| 18 | +#include "register/register.h" | ||
| 19 | +#include "tiling/platform/platform_ascendc.h" | ||
| 20 | +#include "ops_log.h" | ||
| 21 | +#include <algorithm> | ||
| 22 | + | ||
| 23 | +constexpr uint32_t RESERVED_UB_SIZE = 1024; | ||
| 24 | +constexpr uint32_t WORKSPACE_SIZE = 4 * 1024 * 1024; | ||
| 25 | +constexpr uint32_t ALIGN = 32; | ||
| 26 | +constexpr uint32_t MAX_TENSOR_SIZE = 24 * 1024; // 当Tensor长度大于此值时调整分核策略性能更优 | ||
| 27 | + | ||
| 28 | +namespace optiling { | ||
| 29 | + | ||
| 30 | +static ge::graphStatus TilingFunc(gert::TilingContext* context) | ||
| 31 | +{ | ||
| 32 | + // 获取硬件信息 | ||
| 33 | + uint64_t ubSize = 0; | ||
| 34 | + const auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); | ||
| 35 | + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); | ||
| 36 | + uint64_t maxUbSize = ubSize - RESERVED_UB_SIZE; | ||
| 37 | + uint64_t coreNum = ascendcPlatform.GetCoreNumAiv(); | ||
| 38 | + if (coreNum == 0) { | ||
| 39 | + OPS_LOG_E(context, "[ERROR] need more than 0 ai core"); | ||
| 40 | + return ge::GRAPH_FAILED; | ||
| 41 | + } | ||
| 42 | + // 计算偏移地址和大小 | ||
| 43 | + auto offsets = context->GetAttrs()->GetListInt(0); | ||
| 44 | + int32_t offsetLen = *context->GetAttrs()->GetInt(1); | ||
| 45 | + int32_t jtNum = *context->GetAttrs()->GetInt(2); | ||
| 46 | + int32_t nPrefixFromRight = *context->GetAttrs()->GetInt(3); | ||
| 47 | + | ||
| 48 | + uint32_t sliceSize[MAX_SLICE_SIZE + 1] = {0}; // 当nPrefixFromRight > 0时,切片个数+1。 | ||
| 49 | + uint32_t inputOffsetBegin[MAX_SLICE_SIZE + 1] = {0}; | ||
| 50 | + uint32_t outputOffsetBegin[MAX_SLICE_SIZE + 1] = {0}; | ||
| 51 | + | ||
| 52 | + uint32_t indexSliceForRight = offsetLen - 1; | ||
| 53 | + uint32_t outputOffset = 0; | ||
| 54 | + const int64_t* offsetsData = offsets->GetData(); | ||
| 55 | + | ||
| 56 | + uint32_t maxTensorLen = std::max<int32_t>(offsetsData[indexSliceForRight], offsetsData[2 * offsetLen - 1]); | ||
| 57 | + | ||
| 58 | + // 判断是否需要优化切片策略：Tensor长度较长且分片较小时 | ||
| 59 | + const bool needAdjustSliceStrategy = (offsetLen < coreNum) && (maxTensorLen > MAX_TENSOR_SIZE); | ||
| 60 | + | ||
| 61 | + if (nPrefixFromRight == 0) { | ||
| 62 | + // [A] + [B] | ||
| 63 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 64 | + // A[:indexSliceForRight-1] | ||
| 65 | + inputOffsetBegin[i] = offsetsData[i]; | ||
| 66 | + sliceSize[i] = offsetsData[i + 1] - offsetsData[i]; | ||
| 67 | + // B[indexSliceForRight:] | ||
| 68 | + inputOffsetBegin[i + indexSliceForRight] = offsetsData[offsetLen + i]; | ||
| 69 | + sliceSize[i + indexSliceForRight] = offsetsData[offsetLen + i + 1] - offsetsData[offsetLen + i]; | ||
| 70 | + } | ||
| 71 | + } else { | ||
| 72 | + // [nPrefixFromB] + [A] + [remainB] | ||
| 73 | + // first offset [nPrefixFromB] | ||
| 74 | + int32_t bLength = offsetsData[offsetLen + 1] - offsetsData[offsetLen]; | ||
| 75 | + // nPrefixFromRight | ||
| 76 | + sliceSize[indexSliceForRight] = std::min<int32_t>(nPrefixFromRight, bLength); | ||
| 77 | + inputOffsetBegin[indexSliceForRight] = 0; | ||
| 78 | + // [A] + ([remainB] + [nPrefixFromB]) | ||
| 79 | + for (int i = 0; i < offsetLen - 2; i++) { | ||
| 80 | + // [A] | ||
| 81 | + inputOffsetBegin[i] = offsetsData[i]; | ||
| 82 | + sliceSize[i] = offsetsData[i + 1] - offsetsData[i]; | ||
| 83 | + // [remainB] | ||
| 84 | + int32_t remainB = std::max<int32_t>((bLength - nPrefixFromRight), 0); | ||
| 85 | + // [next nPrefixFromB] | ||
| 86 | + bLength = offsetsData[offsetLen + i + 2] - offsetsData[offsetLen + i + 1]; | ||
| 87 | + int32_t prefixNextB = std::min<int32_t>(nPrefixFromRight, bLength); | ||
| 88 | + | ||
| 89 | + inputOffsetBegin[i + offsetLen] = inputOffsetBegin[offsetLen + i - 1] + sliceSize[i + offsetLen - 1]; | ||
| 90 | + sliceSize[i + offsetLen] = remainB + prefixNextB; | ||
| 91 | + } | ||
| 92 | + // last offset [A] + [remainB] | ||
| 93 | + inputOffsetBegin[indexSliceForRight - 1] = offsetsData[indexSliceForRight - 1]; | ||
| 94 | + sliceSize[indexSliceForRight - 1] = offsetsData[indexSliceForRight] - offsetsData[indexSliceForRight - 1]; | ||
| 95 | + // B | ||
| 96 | + int32_t LastBLength = offsetsData[2 * offsetLen - 1] - offsetsData[2 * offsetLen - 2]; | ||
| 97 | + int32_t lastRemainB = std::max<int32_t>((LastBLength - nPrefixFromRight), 0); | ||
| 98 | + | ||
| 99 | + inputOffsetBegin[2 * offsetLen - 2] = inputOffsetBegin[2 * offsetLen - 3] + sliceSize[2 * offsetLen - 3]; | ||
| 100 | + sliceSize[2 * offsetLen - 2] = lastRemainB; | ||
| 101 | + } | ||
| 102 | + | ||
| 103 | + uint32_t batchSize = (nPrefixFromRight == 0) ? (jtNum * (offsetLen - 1)) : (jtNum * (offsetLen - 1) + 1); | ||
| 104 | + | ||
| 105 | + if (nPrefixFromRight == 0) { | ||
| 106 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 107 | + outputOffsetBegin[i] = outputOffset; | ||
| 108 | + outputOffset += sliceSize[i]; | ||
| 109 | + outputOffsetBegin[i + offsetLen - 1] = outputOffset; | ||
| 110 | + outputOffset += sliceSize[i + offsetLen - 1]; | ||
| 111 | + } | ||
| 112 | + } else { | ||
| 113 | + outputOffsetBegin[offsetLen - 1] = outputOffset; | ||
| 114 | + outputOffset += sliceSize[offsetLen - 1]; | ||
| 115 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 116 | + outputOffsetBegin[i] = outputOffset; | ||
| 117 | + outputOffset += sliceSize[i]; | ||
| 118 | + outputOffsetBegin[i + offsetLen] = outputOffset; | ||
| 119 | + outputOffset += sliceSize[i + offsetLen]; | ||
| 120 | + } | ||
| 121 | + } | ||
| 122 | + | ||
| 123 | + uint32_t oneSliceSize = maxTensorLen / coreNum; | ||
| 124 | + uint32_t resetIndexSliceForRight = indexSliceForRight; | ||
| 125 | + | ||
| 126 | + std::vector<uint32_t> readjustSliceSize; | ||
| 127 | + std::vector<uint32_t> readjustInputOffsetBegin; | ||
| 128 | + std::vector<uint32_t> readjustOutputOffsetBegin; | ||
| 129 | + | ||
| 130 | + // Tensor长度较长且分片较小时，优化切片策略。 | ||
| 131 | + if (needAdjustSliceStrategy) { | ||
| 132 | + for (int i = 0; i < batchSize; ++i) { | ||
| 133 | + if (i == indexSliceForRight) { | ||
| 134 | + resetIndexSliceForRight = readjustSliceSize.size(); | ||
| 135 | + } | ||
| 136 | + | ||
| 137 | + uint32_t inputSize = sliceSize[i]; | ||
| 138 | + uint32_t inputOffset = inputOffsetBegin[i]; | ||
| 139 | + uint32_t outputOffset = outputOffsetBegin[i]; | ||
| 140 | + | ||
| 141 | + if (inputSize <= oneSliceSize) { | ||
| 142 | + // 不需要切分，直接添加 | ||
| 143 | + readjustSliceSize.push_back(inputSize); | ||
| 144 | + readjustInputOffsetBegin.push_back(inputOffset); | ||
| 145 | + readjustOutputOffsetBegin.push_back(outputOffset); | ||
| 146 | + } else { | ||
| 147 | + // 需要切分成多个分片 | ||
| 148 | + uint32_t remaining = inputSize; | ||
| 149 | + uint32_t currentInputOffset = inputOffset; | ||
| 150 | + uint32_t currentOutputOffset = outputOffset; | ||
| 151 | + | ||
| 152 | + while (remaining > 0) { | ||
| 153 | + uint32_t chunkSize = std::min(remaining, oneSliceSize); | ||
| 154 | + readjustSliceSize.push_back(chunkSize); | ||
| 155 | + readjustInputOffsetBegin.push_back(currentInputOffset); | ||
| 156 | + readjustOutputOffsetBegin.push_back(currentOutputOffset); | ||
| 157 | + currentInputOffset += chunkSize; | ||
| 158 | + currentOutputOffset += chunkSize; | ||
| 159 | + remaining -= chunkSize; | ||
| 160 | + } | ||
| 161 | + } | ||
| 162 | + } | ||
| 163 | + } | ||
| 164 | + | ||
| 165 | + uint64_t resetBatchSize = batchSize; | ||
| 166 | + if (needAdjustSliceStrategy) { | ||
| 167 | + resetBatchSize = readjustSliceSize.size(); | ||
| 168 | + } | ||
| 169 | + | ||
| 170 | + // 计算切分逻辑 | ||
| 171 | + uint64_t formerCore = 0; | ||
| 172 | + uint64_t tailCore = 0; | ||
| 173 | + uint64_t batchNumInTail = 0; | ||
| 174 | + uint64_t batchNumInFormer = 0; | ||
| 175 | + if (resetBatchSize < coreNum) { | ||
| 176 | + formerCore = resetBatchSize; | ||
| 177 | + batchNumInFormer = 1; | ||
| 178 | + } else { | ||
| 179 | + formerCore = resetBatchSize % coreNum; | ||
| 180 | + tailCore = coreNum - formerCore; | ||
| 181 | + batchNumInTail = resetBatchSize / coreNum; | ||
| 182 | + batchNumInFormer = batchNumInTail + 1; | ||
| 183 | + } | ||
| 184 | + | ||
| 185 | + // 计算copy最大行数 | ||
| 186 | + OPS_CHECK_PTR_NULL(context->GetRequiredInputShape(0), return ge::GRAPH_FAILED); | ||
| 187 | + OPS_LOG_E_IF_NULL("inputXDesc", context->GetInputDesc(0), return ge::GRAPH_FAILED); | ||
| 188 | + uint32_t inputColSize = context->GetRequiredInputShape(0)->GetOriginShape().GetDim(1); | ||
| 189 | + auto inputXDesc = context->GetInputDesc(0); | ||
| 190 | + ge::DataType xDtype{ge::DT_FLOAT}; | ||
| 191 | + xDtype = inputXDesc->GetDataType(); | ||
| 192 | + uint32_t ubMaxLength = maxUbSize / ge::GetSizeByDataType(xDtype); | ||
| 193 | + ubMaxLength = ubMaxLength & ~(ALIGN - 1); | ||
| 194 | + | ||
| 195 | + // 设置tiling | ||
| 196 | + ConcatJaggedTensorTilingData tiling; | ||
| 197 | + tiling.set_jtNum(jtNum); | ||
| 198 | + tiling.set_inputColSize(inputColSize); | ||
| 199 | + tiling.set_ubMaxLength(ubMaxLength); | ||
| 200 | + tiling.set_formerCore(formerCore); | ||
| 201 | + tiling.set_tailCore(tailCore); | ||
| 202 | + tiling.set_batchNumInTail(batchNumInTail); | ||
| 203 | + tiling.set_batchNumInFormer(batchNumInFormer); | ||
| 204 | + | ||
| 205 | + if (needAdjustSliceStrategy) { | ||
| 206 | + // 填充剩余元素为 0 | ||
| 207 | + while (readjustSliceSize.size() < MAX_SLICE_SIZE + 1) { | ||
| 208 | + readjustSliceSize.push_back(0); | ||
| 209 | + readjustInputOffsetBegin.push_back(0); | ||
| 210 | + readjustOutputOffsetBegin.push_back(0); | ||
| 211 | + } | ||
| 212 | + tiling.set_indexSliceForRight(resetIndexSliceForRight); | ||
| 213 | + tiling.set_inputOffsetBegin(readjustInputOffsetBegin.data()); | ||
| 214 | + tiling.set_sliceSize(readjustSliceSize.data()); | ||
| 215 | + tiling.set_outputOffsetBegin(readjustOutputOffsetBegin.data()); | ||
| 216 | + | ||
| 217 | + } else { | ||
| 218 | + tiling.set_indexSliceForRight(indexSliceForRight); | ||
| 219 | + tiling.set_inputOffsetBegin(inputOffsetBegin); | ||
| 220 | + tiling.set_sliceSize(sliceSize); | ||
| 221 | + tiling.set_outputOffsetBegin(outputOffsetBegin); | ||
| 222 | + } | ||
| 223 | + | ||
| 224 | + context->SetBlockDim(coreNum); | ||
| 225 | + OPS_CHECK_PTR_NULL(context->GetRawTilingData(), return ge::GRAPH_FAILED); | ||
| 226 | + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); | ||
| 227 | + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); | ||
| 228 | + size_t* currentWorkspace = context->GetWorkspaceSizes(1); | ||
| 229 | + OPS_LOG_E_IF_NULL("currentWorkspace", currentWorkspace, return ge::GRAPH_FAILED); | ||
| 230 | + currentWorkspace[0] = WORKSPACE_SIZE; | ||
| 231 | + | ||
| 232 | + return ge::GRAPH_SUCCESS; | ||
| 233 | +} | ||
| 234 | +} // namespace optiling | ||
| 235 | + | ||
| 236 | +namespace ge { | ||
| 237 | +static ge::graphStatus InferShape(gert::InferShapeContext* context) | ||
| 238 | +{ | ||
| 239 | + auto offsets = context->GetAttrs()->GetListInt(0); | ||
| 240 | + int32_t offsetLen = *context->GetAttrs()->GetInt(1); | ||
| 241 | + int32_t jtNum = *context->GetAttrs()->GetInt(2); | ||
| 242 | + int64_t outputRow = 0; | ||
| 243 | + for (int i = 0; i < jtNum; i++) { | ||
| 244 | + outputRow += offsets->GetData()[i * offsetLen]; | ||
| 245 | + } | ||
| 246 | + const gert::Shape* x_shape = context->GetInputShape(0); | ||
| 247 | + if (x_shape == nullptr) { | ||
| 248 | + OPS_LOG_E("[ERROR]", "InputShape should not be nullptr."); | ||
| 249 | + return ge::GRAPH_FAILED; | ||
| 250 | + } | ||
| 251 | + uint32_t dimNum = x_shape->GetDimNum(); | ||
| 252 | + gert::Shape* y_shape = context->GetOutputShape(0); | ||
| 253 | + if (y_shape == nullptr) { | ||
| 254 | + OPS_LOG_E("[ERROR]", "OutputShape should not be nullptr."); | ||
| 255 | + return ge::GRAPH_FAILED; | ||
| 256 | + } | ||
| 257 | + y_shape->SetDimNum(dimNum); | ||
| 258 | + y_shape->SetDim(0, outputRow); | ||
| 259 | + y_shape->SetDim(1, x_shape->GetDim(1)); | ||
| 260 | + return ge::GRAPH_SUCCESS; | ||
| 261 | +} | ||
| 262 | +static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) | ||
| 263 | +{ | ||
| 264 | + const auto inputDataType = context->GetInputDataType(0); | ||
| 265 | + context->SetOutputDataType(0, inputDataType); | ||
| 266 | + return ge::GRAPH_SUCCESS; | ||
| 267 | +} | ||
| 268 | +} // namespace ge | ||
| 269 | + | ||
| 270 | +namespace ops { | ||
| 271 | +class ConcatJaggedTensor : public OpDef { | ||
| 272 | +public: | ||
| 273 | + explicit ConcatJaggedTensor(const char* name) : OpDef(name) | ||
| 274 | + { | ||
| 275 | + this->Input("values") | ||
| 276 | + .ParamType(DYNAMIC) | ||
| 277 | + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32}) | ||
| 278 | + .FormatList({ge::FORMAT_ND}) | ||
| 279 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 280 | + this->Output("result") | ||
| 281 | + .ParamType(REQUIRED) | ||
| 282 | + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32}) | ||
| 283 | + .FormatList({ge::FORMAT_ND}) | ||
| 284 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 285 | + | ||
| 286 | + this->Attr("offsets").ListInt(); | ||
| 287 | + this->Attr("offsetLen").Int(0); | ||
| 288 | + this->Attr("jtNum").Int(0); | ||
| 289 | + this->Attr("nPrefixFromRight").Int(0); | ||
| 290 | + | ||
| 291 | + this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); | ||
| 292 | + | ||
| 293 | + this->AICore().SetTiling(optiling::TilingFunc); | ||
| 294 | + this->AICore().AddConfig("ascend910b"); | ||
| 295 | + this->AICore().AddConfig("ascend910_93"); | ||
| 296 | +#ifdef SUPPORT_950 | ||
| 297 | + this->AICore().AddConfig("ascend950"); | ||
| 298 | +#endif | ||
| 299 | + } | ||
| 300 | +}; | ||
| 301 | +OP_ADD(ConcatJaggedTensor); | ||
| 302 | +} // namespace ops | ||
| @@ -0,0 +1,39 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | +#ifndef CONCAT_JAGGED_TENSOR_TILING_H | ||
| 16 | +#define CONCAT_JAGGED_TENSOR_TILING_H | ||
| 17 | + | ||
| 18 | +#include "register/tilingdata_base.h" | ||
| 19 | + | ||
| 20 | +constexpr uint32_t MAX_SLICE_SIZE = 2 * 1024; | ||
| 21 | +namespace optiling { | ||
| 22 | +BEGIN_TILING_DATA_DEF(ConcatJaggedTensorTilingData) | ||
| 23 | +TILING_DATA_FIELD_DEF(uint32_t, jtNum); | ||
| 24 | +TILING_DATA_FIELD_DEF(uint32_t, inputColSize); | ||
| 25 | +TILING_DATA_FIELD_DEF(uint32_t, ubMaxLength); | ||
| 26 | +TILING_DATA_FIELD_DEF(uint32_t, formerCore); | ||
| 27 | +TILING_DATA_FIELD_DEF(uint32_t, tailCore); | ||
| 28 | +TILING_DATA_FIELD_DEF(uint32_t, batchNumInTail); | ||
| 29 | +TILING_DATA_FIELD_DEF(uint32_t, batchNumInFormer); | ||
| 30 | +TILING_DATA_FIELD_DEF(uint32_t, indexSliceForRight); | ||
| 31 | + | ||
| 32 | +TILING_DATA_FIELD_DEF_ARR(uint32_t, MAX_SLICE_SIZE + 1, inputOffsetBegin) | ||
| 33 | +TILING_DATA_FIELD_DEF_ARR(uint32_t, MAX_SLICE_SIZE + 1, sliceSize) | ||
| 34 | +TILING_DATA_FIELD_DEF_ARR(uint32_t, MAX_SLICE_SIZE + 1, outputOffsetBegin) | ||
| 35 | +END_TILING_DATA_DEF; | ||
| 36 | + | ||
| 37 | +REGISTER_TILING_DATA_CLASS(ConcatJaggedTensor, ConcatJaggedTensorTilingData) | ||
| 38 | +} // namespace optiling | ||
| 39 | +#endif | ||
| @@ -0,0 +1,24 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# concat_2d_jagged/op_kernel：AscendC kernel 编译 | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 使用 CANN 的 npu_op_kernel_* 函数编译 kernel。 | ||
| 5 | +# npu_op_package 已在上层 CMakeLists.txt 调用。 | ||
| 6 | +# OPS_REC_SOC_CXX_DEFS 由根 CMakeLists.txt 统一设置（SUPPORT_950 / SUPPORT_A2 / SUPPORT_A3） | ||
| 7 | +set(_KERNEL_TARGET "${VENDOR_NAME}_kernels") | ||
| 8 | + | ||
| 9 | +npu_op_kernel_options(${_KERNEL_TARGET} ALL | ||
| 10 | + OPTIONS "--cce-long-call=true" ${OPS_REC_SOC_CXX_DEFS} | ||
| 11 | +) | ||
| 12 | + | ||
| 13 | +npu_op_kernel_sources(${_KERNEL_TARGET} | ||
| 14 | + KERNEL_DIR ./ | ||
| 15 | +) | ||
| 16 | + | ||
| 17 | +npu_op_kernel_library(${_KERNEL_TARGET} | ||
| 18 | + SRC_BASE ${CMAKE_CURRENT_SOURCE_DIR}/ | ||
| 19 | + TILING_LIBRARY cust_optiling | ||
| 20 | +) | ||
| 21 | + | ||
| 22 | +npu_op_package_add(${VENDOR_NAME} | ||
| 23 | + LIBRARY ${_KERNEL_TARGET} | ||
| 24 | +) | ||
| @@ -0,0 +1,27 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +#include "kernel_operator.h" | ||
| 17 | +#include "concat_jagged_tensor.h" | ||
| 18 | + | ||
| 19 | +using namespace AscendC; | ||
| 20 | + | ||
| 21 | +extern "C" __global__ __aicore__ void concat_jagged_tensor(GM_ADDR values, GM_ADDR result, | ||
| 22 | + GM_ADDR workspace, GM_ADDR tiling) { | ||
| 23 | + GET_TILING_DATA(tiling_data, tiling); | ||
| 24 | + ConcatJaggedTensorKernel<DTYPE_VALUES> op; | ||
| 25 | + op.Init(values, result, workspace, &tiling_data); | ||
| 26 | + op.Process(values, result, workspace, &tiling_data); | ||
| 27 | +} | ||
| @@ -0,0 +1,124 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +#ifndef CONCAT_JAGGED_TENSOR_H | ||
| 17 | +#define CONCAT_JAGGED_TENSOR_H | ||
| 18 | +#include "kernel_operator.h" | ||
| 19 | + | ||
| 20 | +using namespace AscendC; | ||
| 21 | + | ||
| 22 | +#define HardEvent_SetWait_Flag(event_name, SRC, DST) \ | ||
| 23 | + { \ | ||
| 24 | + event_t event_name = static_cast<event_t>(GetTPipePtr()->FetchEventID(AscendC::HardEvent::SRC##_##DST)); \ | ||
| 25 | + set_flag(PIPE_##SRC, PIPE_##DST, event_name); \ | ||
| 26 | + wait_flag(PIPE_##SRC, PIPE_##DST, event_name); \ | ||
| 27 | + } | ||
| 28 | + | ||
| 29 | +namespace AscendC { | ||
| 30 | +template <typename T> | ||
| 31 | +class ConcatJaggedTensorKernel { | ||
| 32 | +public: | ||
| 33 | + __aicore__ inline ConcatJaggedTensorKernel() = default; | ||
| 34 | + | ||
| 35 | + __aicore__ inline void Init(GM_ADDR values, GM_ADDR result, GM_ADDR workspace, | ||
| 36 | + ConcatJaggedTensorTilingData* tilingData) | ||
| 37 | + { | ||
| 38 | + blockIdx_ = GetBlockIdx(); | ||
| 39 | + jtNum_ = tilingData->jtNum; | ||
| 40 | + inputColSize_ = tilingData->inputColSize; | ||
| 41 | + ubMaxLength_ = tilingData->ubMaxLength; | ||
| 42 | + indexSliceForRight_ = tilingData->indexSliceForRight; | ||
| 43 | + | ||
| 44 | + formerCore_ = tilingData->formerCore; | ||
| 45 | + tailCore_ = tilingData->tailCore; | ||
| 46 | + batchNumInTail_ = tilingData->batchNumInTail; | ||
| 47 | + batchNumInFormer_ = tilingData->batchNumInFormer; | ||
| 48 | + | ||
| 49 | + batchNum_ = blockIdx_ < formerCore_ ? tilingData->batchNumInFormer : tilingData->batchNumInTail; | ||
| 50 | + if (blockIdx_ < formerCore_) { | ||
| 51 | + blockBatchOffsetBase_ = blockIdx_ * tilingData->batchNumInFormer; | ||
| 52 | + } else { | ||
| 53 | + blockBatchOffsetBase_ = | ||
| 54 | + formerCore_ * tilingData->batchNumInFormer + (blockIdx_ - formerCore_) * tilingData->batchNumInTail; | ||
| 55 | + } | ||
| 56 | + | ||
| 57 | + pipe_.InitBuffer(moveQue_, ubMaxLength_ * sizeof(T)); | ||
| 58 | + } | ||
| 59 | + | ||
| 60 | + __aicore__ inline void Process(GM_ADDR values, GM_ADDR result, GM_ADDR workspace, | ||
| 61 | + ConcatJaggedTensorTilingData* tilingData) | ||
| 62 | + { | ||
| 63 | + for (uint32_t i = 0; i < batchNum_; i++) { | ||
| 64 | + int64_t index = blockBatchOffsetBase_ + i; | ||
| 65 | + int64_t output_offset_base = tilingData->outputOffsetBegin[index] * inputColSize_; | ||
| 66 | + int64_t slice_size = tilingData->sliceSize[index] * inputColSize_; | ||
| 67 | + int64_t tensor_index = (index < indexSliceForRight_) ? 0 : 1; | ||
| 68 | + | ||
| 69 | + int64_t input_offset = 0; | ||
| 70 | + | ||
| 71 | + inputGm_.SetGlobalBuffer((__gm__ T*)GetTensorAddr(values, tensor_index) + | ||
| 72 | + tilingData->inputOffsetBegin[index] * inputColSize_); | ||
| 73 | + outputGm_.SetGlobalBuffer((__gm__ T*)result + output_offset_base); | ||
| 74 | + while (slice_size > 0) { | ||
| 75 | + int64_t copySize = slice_size > ubMaxLength_ ? ubMaxLength_ : slice_size; | ||
| 76 | + | ||
| 77 | + LocalTensor<T> moveTensor = moveQue_.AllocTensor<T>(); | ||
| 78 | + DataCopyExtParams copyParams{(uint16_t)1, static_cast<uint32_t>(copySize * sizeof(T)), 0, 0, 0}; | ||
| 79 | + DataCopyPad(moveTensor, inputGm_[input_offset], copyParams, padParams); | ||
| 80 | + HardEvent_SetWait_Flag(xLocal_copyin, MTE2, MTE3); | ||
| 81 | + | ||
| 82 | + DataCopyPad(outputGm_[input_offset], moveTensor, copyParams); | ||
| 83 | + slice_size -= copySize; | ||
| 84 | + input_offset += copySize; | ||
| 85 | + HardEvent_SetWait_Flag(xLocal_copyout, MTE3, MTE2); | ||
| 86 | + | ||
| 87 | + moveQue_.FreeTensor(moveTensor); | ||
| 88 | + } | ||
| 89 | + } | ||
| 90 | + } | ||
| 91 | + | ||
| 92 | +private: | ||
| 93 | + __aicore__ inline __gm__ T* GetTensorAddr(GM_ADDR tensorList, uint32_t index) | ||
| 94 | + { | ||
| 95 | + __gm__ uint64_t* dataAddr = reinterpret_cast<__gm__ uint64_t*>(tensorList); | ||
| 96 | + uint64_t tensorPtrOffset = *dataAddr; | ||
| 97 | + | ||
| 98 | + __gm__ uint64_t* tensorPtr = dataAddr + (tensorPtrOffset >> 3); | ||
| 99 | + return reinterpret_cast<__gm__ T*>(*(tensorPtr + index)); | ||
| 100 | + } | ||
| 101 | + | ||
| 102 | +private: | ||
| 103 | + TPipe pipe_; | ||
| 104 | + GlobalTensor<T> inputGm_; | ||
| 105 | + GlobalTensor<T> outputGm_; | ||
| 106 | + TBuf<QuePosition::VECCALC> moveQue_; | ||
| 107 | + | ||
| 108 | + uint32_t blockIdx_; | ||
| 109 | + uint32_t jtNum_; | ||
| 110 | + uint32_t inputColSize_; | ||
| 111 | + uint32_t ubMaxLength_; | ||
| 112 | + uint32_t formerCore_; | ||
| 113 | + uint32_t tailCore_; | ||
| 114 | + uint32_t batchNumInTail_; | ||
| 115 | + uint32_t batchNumInFormer_; | ||
| 116 | + uint32_t indexSliceForRight_; | ||
| 117 | + | ||
| 118 | + uint32_t batchNum_; | ||
| 119 | + uint32_t blockBatchOffsetBase_; | ||
| 120 | + | ||
| 121 | + DataCopyPadExtParams<T> padParams{true, 0, 0, 0}; | ||
| 122 | +}; | ||
| 123 | +} // namespace AscendC | ||
| 124 | +#endif | ||
| @@ -0,0 +1,517 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# -*- coding: utf-8 -*- | ||
| 3 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 4 | +# | ||
| 5 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 6 | +# you may not use this file except in compliance with the License. | ||
| 7 | +# You may obtain a copy of the License at | ||
| 8 | +# | ||
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 10 | +# | ||
| 11 | +# Unless required by applicable law or agreed to in writing, software | ||
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 14 | +# See the License for the specific language governing permissions and | ||
| 15 | +# limitations under the License. | ||
| 16 | +# ============================================================================== | ||
| 17 | + | ||
| 18 | +""" | ||
| 19 | +concat_2d_jagged 算子泛化测试 | ||
| 20 | + | ||
| 21 | +覆盖: | ||
| 22 | + - 前向正确性（多 shape / 多 col / 多 dtype） | ||
| 23 | + - nPrefixFromRight 参数（多取值） | ||
| 24 | + - 反向自动求导（concat.backward ≡ split） | ||
| 25 | + - torch.compile / dynamo graph mode | ||
| 26 | + - 参数校验异常分支 | ||
| 27 | +""" | ||
| 28 | + | ||
| 29 | +import random | ||
| 30 | +import sys | ||
| 31 | + | ||
| 32 | +import pytest | ||
| 33 | +import numpy as np | ||
| 34 | +import torch | ||
| 35 | + | ||
| 36 | +import ops_rec | ||
| 37 | + | ||
| 38 | +DEVICE = "npu:0" | ||
| 39 | + | ||
| 40 | +debug = False | ||
| 41 | +if debug: | ||
| 42 | + seed = 42 | ||
| 43 | + random.seed(seed) | ||
| 44 | + | ||
| 45 | +torch.npu.config.allow_internal_format = False | ||
| 46 | +torch.npu.set_device(DEVICE) | ||
| 47 | + | ||
| 48 | +# 测试 shape 组合较多，增大 cache 避免 dynamic=True 时因重编译超限报错 | ||
| 49 | +torch._dynamo.config.cache_size_limit = 256 | ||
| 50 | + | ||
| 51 | +# ---------- 误差范围 ---------- | ||
| 52 | +_PRECISION_ERROR_RANGE = { | ||
| 53 | + torch.float32: 1e-4, | ||
| 54 | + torch.float16: 1e-4, | ||
| 55 | + torch.bfloat16: 1e-4, | ||
| 56 | + torch.int32: 0, # int32 为纯搬运，应精确 | ||
| 57 | +} | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +def concat_jagged_tensor_golden(value_a, value_b, offset_a, offset_b, max_seqlens): | ||
| 61 | + """nPrefixFromRight=0 的纯 PyTorch 参考实现：按 offset 交错拼接""" | ||
| 62 | + valuesList = [value_a, value_b] | ||
| 63 | + offsetList = [offset_a, offset_b] | ||
| 64 | + output = [] | ||
| 65 | + offsetListLen = len(offsetList) | ||
| 66 | + offsetLen = len(offsetList[0]) | ||
| 67 | + for j in range(offsetLen - 1): | ||
| 68 | + for i in range(offsetListLen): | ||
| 69 | + for k in range(offsetList[i][j], offsetList[i][j + 1]): | ||
| 70 | + output.append(valuesList[i][k]) | ||
| 71 | + return torch.stack(output) | ||
| 72 | + | ||
| 73 | + | ||
| 74 | +def concat_jagged_tensor_golden_with_prefix(value_a, value_b, offset_a, offset_b, n_prefix_from_right): | ||
| 75 | + """ | ||
| 76 | + 参考实现 for concat with nPrefixFromRight parameter. | ||
| 77 | + Output structure: [B_prefix] + [A] + [B_remaining] | ||
| 78 | + """ | ||
| 79 | + offsetList = [offset_a, offset_b] | ||
| 80 | + offsetLen = len(offsetList[0]) | ||
| 81 | + output = [] | ||
| 82 | + | ||
| 83 | + for j in range(offsetLen - 1): | ||
| 84 | + a_start, a_end = offsetList[0][j], offsetList[0][j + 1] | ||
| 85 | + b_start, b_end = offsetList[1][j], offsetList[1][j + 1] | ||
| 86 | + | ||
| 87 | + b_len = b_end - b_start | ||
| 88 | + prefix_size = min(n_prefix_from_right, b_len) | ||
| 89 | + | ||
| 90 | + # B prefix: first n elements from B | ||
| 91 | + for k in range(b_start, min(b_start + prefix_size, b_end)): | ||
| 92 | + output.append(value_b[k]) | ||
| 93 | + | ||
| 94 | + # A: all elements from A | ||
| 95 | + for k in range(a_start, a_end): | ||
| 96 | + output.append(value_a[k]) | ||
| 97 | + | ||
| 98 | + # B remaining: remaining elements from B | ||
| 99 | + for k in range(b_start + prefix_size, b_end): | ||
| 100 | + output.append(value_b[k]) | ||
| @@ -0,0 +1,27 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# split_2d_jagged 算子聚合 CMakeLists | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 设置 vendor 名，使用 npu_op_package + INSTALL_PATH 将产物 | ||
| 5 | +# 定向到 opp/<variant>/vendors/<op_name>/ 目录。 | ||
| 6 | +# ============================================================================ | ||
| 7 | + | ||
| 8 | +# vendor 名，同时也是 CANN package 名 | ||
| 9 | +set(VENDOR_NAME "split_2d_jagged") | ||
| 10 | + | ||
| 11 | +# 独立 autogen 路径，避免多个算子共享冲突 | ||
| 12 | +set(ASCEND_AUTOGEN_PATH "${CMAKE_CURRENT_BINARY_DIR}/autogen") | ||
| 13 | +file(MAKE_DIRECTORY "${ASCEND_AUTOGEN_PATH}") | ||
| 14 | + | ||
| 15 | +# 安装路径 | ||
| 16 | +set(OPP_VENDOR_INSTALL_PREFIX | ||
| 17 | + "${CUSTOM_OPP_INSTALL_PREFIX}/vendors/${VENDOR_NAME}") | ||
| 18 | + | ||
| 19 | +npu_op_package(${VENDOR_NAME} | ||
| 20 | + TYPE RUN | ||
| 21 | + CONFIG | ||
| 22 | + INSTALL_PATH ${OPP_VENDOR_INSTALL_PREFIX} | ||
| 23 | + VENDOR_NAME ${VENDOR_NAME} | ||
| 24 | +) | ||
| 25 | + | ||
| 26 | +add_subdirectory(op_host) | ||
| 27 | +add_subdirectory(op_kernel) | ||
| @@ -0,0 +1,108 @@ | |||
| 1 | +# split_2d_jagged | ||
| 2 | + | ||
| 3 | +本算子仅支持 NPU 调用。 | ||
| 4 | + | ||
| 5 | +## 产品支持情况 | ||
| 6 | + | ||
| 7 | +| 硬件型号 | 是否支持 | | ||
| 8 | +| -------------------- |------| | ||
| 9 | +| Atlas A2训练系列产品 | 是 | | ||
| 10 | +| Atlas A3训练系列产品 | 是 | | ||
| 11 | +| Atlas 推理系列产品 | 否 | | ||
| 12 | + | ||
| 13 | +## 目录结构 | ||
| 14 | + | ||
| 15 | +`split_2d_jagged`（反向）与 `concat_2d_jagged`（前向）同属于 `concat_jagged_tensor` 聚合目录，torch 插件由两个算子共用： | ||
| 16 | + | ||
| 17 | +```shell | ||
| 18 | +concat_jagged_tensor | ||
| 19 | +├── torch_extension # torch 插件（schema 注册、参数校验、autograd 注册，两算子共用） | ||
| 20 | +│ ├── concat_jagged_tensor_plugin.cpp | ||
| 21 | +│ └── concat_jagged_tensor_autograd.py | ||
| 22 | +├── concat_2d_jagged # 前向算子，按 offset 拼接两个 jagged tensor | ||
| 23 | +│ └── （目录结构同下） | ||
| 24 | +└── split_2d_jagged # 反向算子（本目录） | ||
| 25 | + ├── op_host # 算子 host 侧实现（OpDef、tiling） | ||
| 26 | + ├── op_kernel # 算子 kernel 侧实现 | ||
| 27 | + ├── examples # 调用示例（aclnn / torch） | ||
| 28 | + ├── test # pytest 功能测试 | ||
| 29 | + ├── docs # 接口文档 | ||
| 30 | + └── README.md # 算子说明文档（本文档） | ||
| 31 | +``` | ||
| 32 | + | ||
| 33 | +## 功能 | ||
| 34 | + | ||
| 35 | +将一个 tensor 按照 offset 切分成两个不一定等长的 tensor，是 [concat_2d_jagged](../concat_2d_jagged/README.md) 的逆操作。 | ||
| 36 | + | ||
| 37 | +对应接口： | ||
| 38 | + | ||
| 39 | +- aclnn 接口：`aclnnSplitJaggedTensor` | ||
| 40 | +- torch 接口：`ops_rec.ascendc.sequence.split_2d_jagged` | ||
| 41 | + | ||
| 42 | +> CANN 算子名为 `SplitJaggedTensor`，torch 算子名为 `ops_rec::split_2d_jagged`。 | ||
| 43 | + | ||
| 44 | +## 算子实现原理 | ||
| 45 | + | ||
| 46 | +输入: | ||
| 47 | + | ||
| 48 | +```python | ||
| 49 | +values = tensor([[1,1,1,1,1,1,1,1], | ||
| 50 | + [4,4,4,4,4,4,4,4], | ||
| 51 | + [2,2,2,2,2,2,2,2], | ||
| 52 | + [5,5,5,5,5,5,5,5], | ||
| 53 | + [6,6,6,6,6,6,6,6], | ||
| 54 | + [3,3,3,3,3,3,3,3], | ||
| 55 | + [7,7,7,7,7,7,7,7] | ||
| 56 | + ]) | ||
| 57 | +# 长度为2N,前N个数表示第一个tensor的offset,后N个数表示第2个tensor的offset。 | ||
| 58 | +offsets = [0, 1, 2, 3, 0, 1, 3, 4] | ||
| 59 | +# offset的长度，两个tensor的offset长度必须一致。 | ||
| 60 | +offsetLen = 4 | ||
| 61 | +# 切分出的tensor个数，这里仅支持两个。 | ||
| 62 | +jtNum = 2 | ||
| 63 | +``` | ||
| 64 | + | ||
| 65 | +输出： | ||
| 66 | + | ||
| 67 | +```python | ||
| 68 | +result = [tensor([[1,1,1,1,1,1,1,1], | ||
| 69 | + [2,2,2,2,2,2,2,2], | ||
| 70 | + [3,3,3,3,3,3,3,3] | ||
| 71 | + ]), | ||
| 72 | + tensor([[4,4,4,4,4,4,4,4], | ||
| 73 | + [5,5,5,5,5,5,5,5], | ||
| 74 | + [6,6,6,6,6,6,6,6], | ||
| 75 | + [7,7,7,7,7,7,7,7], | ||
| 76 | + ])] | ||
| 77 | +``` | ||
| 78 | + | ||
| 79 | +`nPrefixToRight > 0` 时，按"第 2 个输出每个分片的前 `nPrefixToRight` 行已被前移"的排布还原切分，即 `concat_2d_jagged` 对应 `nPrefixFromRight` 参数的逆操作。 | ||
| 80 | + | ||
| 81 | +算子采用 AscendC 开发。host 侧 tiling 把一次切分拆解为 `jtNum * (offsetLen - 1)` 个连续行切片（`nPrefixToRight > 0` 时加 1），均摊到全部 AI Vector 核（前 formerCore 个核各多承担 1 个切片）；当切片数少于核数且输入行数超过 24K 时，再按核数将长切片等分拆细，均衡各核负载。kernel 侧每个核将自己负责的切片经 UB 中转（`DataCopyPad` 按 ubMaxLength 分段搬运），从输入 GM 搬写到两个输出 GM 的对应位置。 | ||
| 82 | + | ||
| 83 | +## 算子输入与输出 | ||
| 84 | + | ||
| 85 | +| 名称 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | 范围 | 说明 | | ||
| 86 | +|----------------|--------|---------|-----------------------------------|---------------------------|-------------------------------------|-------------------| | ||
| 87 | +| values | 输入 | Tensor | bfloat16/float16/float32/int32 | [dim0, dim1] | 二维tensor | 待切分的tensor | | ||
| 88 | +| offsets | 输入 | List | int64 | List[int] | 长度为2N，前N个为第1个tensor的offset，后N个为第2个tensor的offset | 切分偏移 | | ||
| 89 | +| offsetLen | 输入 | int | int64 | NA | [2, 1024] | 单个tensor的offset长度 | | ||
| 90 | +| jtNum | 输入 | int | int64 | NA | 支持jtNum = 2 | 切分后tensor的个数 | | ||
| 91 | +| nPrefixToRight | 输入 | int | int64 | NA | nPrefixToRight >= 0 | 切分时右侧tensor每个分片被前移的行数 | | ||
| 92 | +| result | 输出 | List | bfloat16/float16/float32/int32 | List[Tensor] | 结果为Tensor列表，且列表长度为2，表示切分后的两个Tensor；两个tensor的dim0分别等于两个offset的末元素，dim1与输入一致 | NA | | ||
| 93 | + | ||
| 94 | +## 约束说明 | ||
| 95 | + | ||
| 96 | +- values 必须为 2D，数据类型仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 97 | +- 每个 tensor 的 offset 需单调不减、首元素为 0，末元素决定对应输出 tensor 的 dim0 | ||
| 98 | +- offsetLen 取值范围 [2, 1024]，jtNum 仅支持 2，nPrefixToRight ≥ 0 | ||
| 99 | +- 该算子需要在 NPU 设备上运行 | ||
| 100 | + | ||
| 101 | +> torch 接口的入参形式（offset_a / offset_b 以 tensor 传入）、autograd 与 torch.compile 支持及扩展层额外校验，详见 [torch 接口文档](./docs/torch_split_2d_jagged.md)。 | ||
| 102 | + | ||
| 103 | +## 调用示例 | ||
| 104 | + | ||
| 105 | +- [aclnn 调用示例](./examples/aclnn/test_aclnn_split_2d_jagged.cpp) | ||
| 106 | +- [torch 调用示例](./examples/torch/test_torch_split_2d_jagged.py) | ||
| 107 | +- [aclnn 接口文档](./docs/aclnn_split_jagged_tensor.md) | ||
| 108 | +- [torch 接口文档](./docs/torch_split_2d_jagged.md) | ||
| @@ -0,0 +1,91 @@ | |||
| 1 | +# aclnnConcatJaggedTensorGrad | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 硬件型号 | 是否支持 | | ||
| 6 | +| -------------------- |------| | ||
| 7 | +| Atlas A2训练系列产品 | 是 | | ||
| 8 | +| Atlas A3训练系列产品 | 是 | | ||
| 9 | +| Atlas 推理系列产品 | 否 | | ||
| 10 | + | ||
| 11 | +## 功能说明 | ||
| 12 | + | ||
| 13 | +将一个 tensor 按照 offset 切分成两个不一定等长的 tensor，是 `aclnnConcatJaggedTensor` 的逆操作。对应 torch 接口 `ops_rec.ascendc.sequence.split_2d_jagged`，CANN 算子名为 `ConcatJaggedTensorGrad`。 | ||
| 14 | + | ||
| 15 | +offsets 以 host 侧 int64 数组属性传入，长度为 `offsetLen` 的 2 倍，前 `offsetLen` 个为第 1 个输出的 offset，后 `offsetLen` 个为第 2 个输出的 offset。两个输出以动态 tensor 列表（`aclTensorList`）形式传出。 | ||
| 16 | + | ||
| 17 | +设 `N = offsetLen`，第 i 个分片中输出 A 的行数为 `offsets[i+1] - offsets[i]`、输出 B 的行数为 `offsets[N+i+1] - offsets[N+i]`： | ||
| 18 | + | ||
| 19 | +- `nPrefixToRight = 0` 时，输入按 `[A_0, B_0, A_1, B_1, ..., A_{N-2}, B_{N-2}]` 的交错排布依次切开，还原出 A、B | ||
| 20 | +- `nPrefixToRight > 0` 时，按"B 每个分片的前 `nPrefixToRight` 行已被前移"的排布切分，即 `aclnnConcatJaggedTensor` 对应 `nPrefixFromRight` 参数的逆操作 | ||
| 21 | + | ||
| 22 | +## 函数原型 | ||
| 23 | + | ||
| 24 | +两段式接口，必须先调用 `aclnnConcatJaggedTensorGradGetWorkspaceSize` 接口获取 workspace 大小和执行器，再调用 `aclnnConcatJaggedTensorGrad` 接口执行计算。 | ||
| 25 | + | ||
| 26 | +```c | ||
| 27 | +aclnnStatus aclnnConcatJaggedTensorGradGetWorkspaceSize( | ||
| 28 | + const aclTensor *values, | ||
| 29 | + const aclIntArray *offsets, | ||
| 30 | + int64_t offsetLen, | ||
| 31 | + int64_t jtNum, | ||
| 32 | + int64_t nPrefixToRight, | ||
| 33 | + const aclTensorList *out, | ||
| 34 | + uint64_t *workspaceSize, | ||
| 35 | + aclOpExecutor **executor); | ||
| 36 | + | ||
| 37 | +aclnnStatus aclnnConcatJaggedTensorGrad( | ||
| 38 | + void *workspace, | ||
| 39 | + uint64_t workspaceSize, | ||
| 40 | + aclOpExecutor *executor, | ||
| 41 | + aclrtStream stream); | ||
| 42 | +``` | ||
| 43 | + | ||
| 44 | +## aclnnConcatJaggedTensorGradGetWorkspaceSize | ||
| 45 | + | ||
| 46 | +### 参数说明 | ||
| 47 | + | ||
| 48 | +| 参数名 | 输入/输出 | 描述 | 数据类型 | 数据格式 | | ||
| 49 | +|--------|---------|------|---------|---------| | ||
| 50 | +| values（aclTensor） | 输入 | 待切分的 tensor，shape [rows, dim1] | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | | ||
| 51 | +| offsets（aclIntArray） | 输入 | host 侧偏移数组，长度为 `2 * offsetLen`，前 offsetLen 个为第 1 个输出的 offset，后 offsetLen 个为第 2 个输出的 offset | INT64 | - | | ||
| 52 | +| offsetLen（int64_t） | 输入 | 单个 tensor 的 offset 长度，取值范围 [2, 1024] | - | - | | ||
| 53 | +| jtNum（int64_t） | 输入 | 切分出的 tensor 个数，仅支持 2 | - | - | | ||
| 54 | +| nPrefixToRight（int64_t） | 输入 | 切分时右侧 tensor 每个分片被前移的行数，取值 ≥ 0 | - | - | | ||
| 55 | +| out（aclTensorList） | 输出 | 切分结果，2 个 tensor 组成的列表：out[0] shape [offsets[offsetLen-1], dim1]，out[1] shape [offsets[2*offsetLen-1], dim1]，数据类型与 values 一致 | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | | ||
| 56 | +| workspaceSize（uint64_t*） | 输出 | 返回需要在 Device 侧申请的 workspace 大小 | - | - | | ||
| 57 | +| executor（aclOpExecutor**） | 输出 | 返回算子执行器 | - | - | | ||
| 58 | + | ||
| 59 | +### 返回值 | ||
| 60 | + | ||
| 61 | +| 返回值 | 错误码 | 描述 | | ||
| 62 | +|--------|-------|------| | ||
| 63 | +| ACLNN_ERR_PARAM_NULLPTR | 161001 | 传入参数是必选输入/输出且为空指针 | | ||
| 64 | +| ACLNN_ERR_PARAM_INVALID | 161002 | 参数数据类型或数据格式不在支持范围内 | | ||
| 65 | + | ||
| 66 | +## aclnnConcatJaggedTensorGrad | ||
| 67 | + | ||
| 68 | +### 参数说明 | ||
| 69 | + | ||
| 70 | +| 参数名 | 输入/输出 | 描述 | | ||
| 71 | +|--------|---------|------| | ||
| 72 | +| workspace | 输入 | 在 Device 侧申请的 workspace 内存地址 | | ||
| 73 | +| workspaceSize | 输入 | workspace 大小，由 `aclnnConcatJaggedTensorGradGetWorkspaceSize` 获取 | | ||
| 74 | +| executor | 输入 | 算子执行器 | | ||
| 75 | +| stream | 输入 | 指定执行任务的 Stream | | ||
| 76 | + | ||
| 77 | +### 返回值 | ||
| 78 | + | ||
| 79 | +aclnnStatus：返回状态码，成功返回 `ACLNN_SUCCESS`。 | ||
| 80 | + | ||
| 81 | +## 约束说明 | ||
| 82 | + | ||
| 83 | +- values 必须为 2D，数据类型仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 84 | +- 每个输出的 offset 需单调不减、首元素为 0，末元素决定对应输出 tensor 的 dim0 | ||
| 85 | +- offsetLen 取值范围 [2, 1024]，jtNum 仅支持 2，nPrefixToRight 取值 ≥ 0 | ||
| 86 | +- out 列表长度必须等于 jtNum（仅支持 2），两个输出 tensor 需由调用侧按 offset 末元素提前分配 | ||
| 87 | +- 该算子需要在 NPU 设备上运行 | ||
| 88 | + | ||
| 89 | +## 调用示例 | ||
| 90 | + | ||
| 91 | +具体参考 [test_aclnn_split_2d_jagged.cpp](../examples/aclnn/test_aclnn_split_2d_jagged.cpp) 的完整实现。 | ||
| @@ -0,0 +1,82 @@ | |||
| 1 | +# split_2d_jagged | ||
| 2 | + | ||
| 3 | +## 产品支持情况 | ||
| 4 | + | ||
| 5 | +| 硬件型号 | 是否支持 | | ||
| 6 | +| -------------------- |------| | ||
| 7 | +| Atlas A2训练系列产品 | 是 | | ||
| 8 | +| Atlas A3训练系列产品 | 是 | | ||
| 9 | +| Atlas 推理系列产品 | 否 | | ||
| 10 | + | ||
| 11 | +## 功能说明 | ||
| 12 | + | ||
| 13 | +将一个 tensor 按照 offset 切分成两个不一定等长的 tensor，是 `ops_rec.ascendc.sequence.concat_2d_jagged` 的逆操作。底层调用 aclnn 接口 `aclnnConcatJaggedTensorGrad`（CANN 算子名 `ConcatJaggedTensorGrad`）。 | ||
| 14 | + | ||
| 15 | +设 `N = offset_a.shape[0]`，第 i 个分片中输出 A 的行数为 `offset_a[i+1] - offset_a[i]`、输出 B 的行数为 `offset_b[i+1] - offset_b[i]`： | ||
| 16 | + | ||
| 17 | +- `n_prefix_to_right = 0` 时，输入按 `[A_0, B_0, A_1, B_1, ..., A_{N-2}, B_{N-2}]` 的交错排布依次切开，还原出 A、B | ||
| 18 | +- `n_prefix_to_right > 0` 时，按"B 每个分片的前 `n_prefix_to_right` 行已被前移"的排布切分，即 `concat_2d_jagged` 对应 `n_prefix_from_right` 参数的逆操作 | ||
| 19 | + | ||
| 20 | +## 接口原型 | ||
| 21 | + | ||
| 22 | +```python | ||
| 23 | +ops_rec.ascendc.sequence.split_2d_jagged( | ||
| 24 | + values: Tensor, | ||
| 25 | + max_seqlen: int, | ||
| 26 | + offset_a: Tensor, | ||
| 27 | + offset_b: Tensor, | ||
| 28 | + dense_size: int = 0, | ||
| 29 | + n_prefix_to_right: int = 0, | ||
| 30 | +) -> Tuple[Tensor, Tensor] | ||
| 31 | +``` | ||
| 32 | + | ||
| 33 | +## 参数说明 | ||
| 34 | + | ||
| 35 | +| 参数名 | 输入/输出 | 参数类型 | 数据类型 | 数据格式 | shape | 说明 | | ||
| 36 | +|--------|---------|---------|---------|---------|-------|------| | ||
| 37 | +| values | 输入 | Tensor | FLOAT32 / FLOAT16 / BF16 / INT32 | ND | [rows, dim1] | 待切分的 tensor | | ||
| 38 | +| max_seqlen | 输入 | int | - | - | - | 保留参数，当前实现未使用 | | ||
| 39 | +| offset_a | 输入 | Tensor | INT64 / INT32 | ND | [N] | 输出 A 的 offset，单调不减、首元素为 0 | | ||
| 40 | +| offset_b | 输入 | Tensor | INT64 / INT32 | ND | [N] | 输出 B 的 offset，长度必须与 offset_a 一致 | | ||
| 41 | +| dense_size | 输入 | int | - | - | - | 保留参数，当前实现未使用 | | ||
| 42 | +| n_prefix_to_right | 输入 | int | - | - | - | 切分时输出 B 每个分片被前移的行数，取值 ≥ 0 | | ||
| 43 | +| output | 输出 | Tuple[Tensor, Tensor] | 与 values 一致 | ND | ([offset_a[-1], dim1], [offset_b[-1], dim1]) | 切分得到的两个 tensor | | ||
| 44 | + | ||
| 45 | +> offset_a / offset_b 以 NPU 上的 tensor 形式传入，插件内部会将其拷贝到 host 并拼接成长度为 2N 的 int64 数组传给底层算子；两个输出的 dim0 分别由 offset_a[-1]、offset_b[-1] 决定。 | ||
| 46 | + | ||
| 47 | +## autograd 与 torch.compile | ||
| 48 | + | ||
| 49 | +- **自动求导**：split 的反向等价于按相同 offset 把两个输出的梯度拼回单个 tensor。由 [concat_jagged_tensor_autograd.py](../../torch_extension/concat_jagged_tensor_autograd.py) 注册，反向内部调用 `ops_rec.ascendc.sequence.concat_2d_jagged` 将 `(grad_a, grad_b)` 拼成 values 的梯度。 | ||
| 50 | +- **torch.compile**：split 输出 shape 依赖 offset tensor 的数据（offset_a[-1] / offset_b[-1]），注册的 fake 实现以 unbacked SymInt 表示两个输出的 dim0；注册时会自动开启 `torch._dynamo.config.capture_dynamic_output_shape_ops = True`，支持 dynamo / inductor 入图（`dynamic=True, fullgraph=True`）。 | ||
| 51 | + | ||
| 52 | +## 约束说明 | ||
| 53 | + | ||
| 54 | +约束分为两层：算子自身语义适用于 aclnn 与 torch 两条调用路径；torch 扩展层额外加的入参校验仅在调用 `ops_rec.ascendc.sequence.split_2d_jagged` 时生效。 | ||
| 55 | + | ||
| 56 | +### 算子约束 | ||
| 57 | + | ||
| 58 | +- values 必须为 2D，数据类型仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 59 | +- offset_a / offset_b 必须为 1D、长度相等，取值范围 [2, 1024]，dtype 为 INT64 / INT32 | ||
| 60 | +- 每个输出的 offset 需单调不减，末元素决定对应输出 tensor 的 dim0 | ||
| 61 | +- 一次仅支持切分出 2 个 tensor（jtNum = 2） | ||
| 62 | +- 该算子需要在 NPU 设备上运行 | ||
| 63 | + | ||
| 64 | +### torch 扩展层额外校验 | ||
| 65 | + | ||
| 66 | +由 [concat_jagged_tensor_plugin.cpp](../../torch_extension/concat_jagged_tensor_plugin.cpp) 在调用算子前完成： | ||
| 67 | + | ||
| 68 | +- 3 个输入 tensor 必须已定义且位于同一 NPU 设备 | ||
| 69 | +- values 必须为 2D，dtype 仅支持 FLOAT32 / FLOAT16 / BF16 / INT32 | ||
| 70 | +- offset_a / offset_b 必须为 1D，dtype 为 INT64 / INT32，长度相等且在 [2, 1024] 范围内 | ||
| 71 | +- `n_prefix_to_right >= 0` | ||
| 72 | +- 需要安装 `torch_npu` | ||
| 73 | + | ||
| 74 | +## 安装方式 | ||
| 75 | + | ||
| 76 | +```bash | ||
| 77 | +bash build.sh --install -j16 | ||
| 78 | +``` | ||
| 79 | + | ||
| 80 | +## 调用示例 | ||
| 81 | + | ||
| 82 | +完整调用示例（含前向、nPrefixToRight、autograd、torch.compile 用例）参考 [test_torch_split_2d_jagged.py](../examples/torch/test_torch_split_2d_jagged.py)。 | ||
Aascendc/sequence/concat_jagged_tensor/split_2d_jagged/examples/aclnn/test_aclnn_split_2d_jagged.cpp+258-0
| @@ -0,0 +1,258 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | +limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +/** | ||
| 17 | + * @file test_aclnn_split_2d_jagged.cpp | ||
| 18 | + * @brief 通过 aclnnSplitJaggedTensor 接口直调 split_2d_jagged 算子 | ||
| 19 | + */ | ||
| 20 | + | ||
| 21 | +#include <cmath> | ||
| 22 | +#include <cstdint> | ||
| 23 | +#include <cstdio> | ||
| 24 | +#include <vector> | ||
| 25 | +#include "acl/acl.h" | ||
| 26 | +#include "aclnn_split_jagged_tensor.h" | ||
| 27 | + | ||
| 28 | +#define CHECK_RET(cond, return_expr) \ | ||
| 29 | + do { \ | ||
| 30 | + if (!(cond)) { \ | ||
| 31 | + return_expr; \ | ||
| 32 | + } \ | ||
| 33 | + } while (0) | ||
| 34 | + | ||
| 35 | +#define LOG_PRINT(message, ...) \ | ||
| 36 | + do { \ | ||
| 37 | + printf(message, ##__VA_ARGS__); \ | ||
| 38 | + } while (0) | ||
| 39 | + | ||
| 40 | +constexpr int64_t JT_NUM = 2; // 一次切分 2 个 tensor | ||
| 41 | + | ||
| 42 | +int64_t GetShapeSize(const std::vector<int64_t>& shape) | ||
| 43 | +{ | ||
| 44 | + int64_t shape_size = 1; | ||
| 45 | + for (auto i : shape) { | ||
| 46 | + shape_size *= i; | ||
| 47 | + } | ||
| 48 | + return shape_size; | ||
| 49 | +} | ||
| 50 | + | ||
| 51 | +int Init(int32_t deviceId, aclrtStream* stream) | ||
| 52 | +{ | ||
| 53 | + auto ret = aclInit(nullptr); | ||
| 54 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); | ||
| 55 | + ret = aclrtSetDevice(deviceId); | ||
| 56 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); | ||
| 57 | + ret = aclrtCreateStream(stream); | ||
| 58 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); | ||
| 59 | + return 0; | ||
| 60 | +} | ||
| 61 | + | ||
| 62 | +template <typename T> | ||
| 63 | +int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, | ||
| 64 | + aclDataType dataType, aclTensor** tensor) | ||
| 65 | +{ | ||
| 66 | + auto size = GetShapeSize(shape) * sizeof(T); | ||
| 67 | + auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 68 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); | ||
| 69 | + ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 70 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); | ||
| 71 | + | ||
| 72 | + std::vector<int64_t> strides(shape.size(), 1); | ||
| 73 | + for (int64_t i = static_cast<int64_t>(shape.size()) - 2; i >= 0; --i) { | ||
| 74 | + strides[i] = shape[i + 1] * strides[i + 1]; | ||
| 75 | + } | ||
| 76 | + | ||
| 77 | + *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, | ||
| 78 | + shape.data(), shape.size(), *deviceAddr); | ||
| 79 | + return 0; | ||
| 80 | +} | ||
| 81 | + | ||
| 82 | +// 执行一次 split 前向：GetWorkspaceSize -> 执行 -> 同步（workspace 在函数内分配与释放） | ||
| 83 | +int RunSplit(aclTensor* values, aclIntArray* offsetsArray, int64_t offsetLen, | ||
| 84 | + int64_t nPrefixToRight, aclTensorList* outList, aclrtStream stream) | ||
| 85 | +{ | ||
| 86 | + uint64_t workspaceSize = 0; | ||
| 87 | + aclOpExecutor* executor = nullptr; | ||
| 88 | + auto ret = aclnnSplitJaggedTensorGetWorkspaceSize(values, offsetsArray, offsetLen, JT_NUM, | ||
| 89 | + nPrefixToRight, outList, | ||
| 90 | + &workspaceSize, &executor); | ||
| 91 | + CHECK_RET(ret == ACL_SUCCESS, | ||
| 92 | + LOG_PRINT("aclnnSplitJaggedTensorGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); | ||
| 93 | + | ||
| 94 | + void* workspaceAddr = nullptr; | ||
| 95 | + if (workspaceSize != 0) { | ||
| 96 | + ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 97 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc workspace failed. ERROR: %d\n", ret); return ret); | ||
| 98 | + } | ||
| 99 | + ret = aclnnSplitJaggedTensor(workspaceAddr, workspaceSize, executor, stream); | ||
| 100 | + if (workspaceAddr != nullptr) { | ||
| 101 | + aclrtFree(workspaceAddr); | ||
| 102 | + } | ||
| 103 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnSplitJaggedTensor failed. ERROR: %d\n", ret); return ret); | ||
| 104 | + | ||
| 105 | + ret = aclrtSynchronizeDevice(); | ||
| 106 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeDevice failed. ERROR: %d\n", ret); return ret); | ||
| 107 | + return 0; | ||
| 108 | +} | ||
| 109 | + | ||
| 110 | +// 逐元素比较，返回错误元素个数 | ||
| 111 | +int64_t CompareBuf(const std::vector<float>& actual, const std::vector<float>& expected, float tol = 1e-6f) | ||
| 112 | +{ | ||
| 113 | + if (actual.size() != expected.size()) { | ||
| 114 | + LOG_PRINT("size mismatch: actual=%zu expected=%zu\n", actual.size(), expected.size()); | ||
| 115 | + return -1; | ||
| 116 | + } | ||
| 117 | + int64_t errCnt = 0; | ||
| 118 | + for (size_t i = 0; i < actual.size(); ++i) { | ||
| 119 | + if (std::fabs(actual[i] - expected[i]) > tol) { | ||
| 120 | + ++errCnt; | ||
| 121 | + } | ||
| 122 | + } | ||
| 123 | + return errCnt; | ||
| 124 | +} | ||
| 125 | + | ||
| 126 | +// 把 {v1, v2, ...} 各扩展为 cols 个元素的一行 | ||
| 127 | +std::vector<float> RepeatRows(const std::vector<float>& rowValues, int64_t cols) | ||
| 128 | +{ | ||
| 129 | + std::vector<float> out; | ||
| 130 | + for (float v : rowValues) { | ||
| 131 | + out.insert(out.end(), cols, v); | ||
| 132 | + } | ||
| 133 | + return out; | ||
| 134 | +} | ||
| 135 | + | ||
| 136 | +// 跑一个 split 用例：创建输入 -> 执行 -> 读回 -> 与 expected 比较，返回错误用例数 | ||
| 137 | +int RunSplitCase(const std::vector<float>& inputHost, int64_t nPrefixToRight, aclIntArray* offsetsArray, | ||
| 138 | + int64_t offsetLen, const char* caseName, aclrtStream stream) | ||
| 139 | +{ | ||
| 140 | + constexpr int64_t kRowsA = 3; | ||
| 141 | + constexpr int64_t kRowsB = 4; | ||
| 142 | + constexpr int64_t kCols = 8; | ||
| 143 | + constexpr int64_t kTotalRows = kRowsA + kRowsB; // 7 | ||
| 144 | + | ||
| 145 | + void* inputDeviceAddr = nullptr; | ||
| 146 | + void* outADeviceAddr = nullptr; | ||
| 147 | + void* outBDeviceAddr = nullptr; | ||
| 148 | + aclTensor* values = nullptr; | ||
| 149 | + aclTensor* outA = nullptr; | ||
| 150 | + aclTensor* outB = nullptr; | ||
| 151 | + | ||
| 152 | + auto ret = CreateAclTensor(inputHost, {kTotalRows, kCols}, &inputDeviceAddr, | ||
| 153 | + aclDataType::ACL_FLOAT, &values); | ||
| 154 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: create input failed\n", caseName); return 1); | ||
| 155 | + ret = CreateAclTensor(std::vector<float>(kRowsA * kCols, 0), {kRowsA, kCols}, &outADeviceAddr, | ||
| 156 | + aclDataType::ACL_FLOAT, &outA); | ||
| 157 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: create outA failed\n", caseName); return 1); | ||
| 158 | + ret = CreateAclTensor(std::vector<float>(kRowsB * kCols, 0), {kRowsB, kCols}, &outBDeviceAddr, | ||
| 159 | + aclDataType::ACL_FLOAT, &outB); | ||
| 160 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: create outB failed\n", caseName); return 1); | ||
| 161 | + | ||
| 162 | + // 动态输出：将 outA、outB 组装为 aclTensorList | ||
| 163 | + aclTensor* outTensors[JT_NUM] = {outA, outB}; | ||
| 164 | + aclTensorList* outList = aclCreateTensorList(outTensors, JT_NUM); | ||
| 165 | + CHECK_RET(outList != nullptr, LOG_PRINT("[FAIL] %s: aclCreateTensorList failed\n", caseName); return 1); | ||
| 166 | + | ||
| 167 | + ret = RunSplit(values, offsetsArray, offsetLen, nPrefixToRight, outList, stream); | ||
| 168 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: run failed\n", caseName); | ||
| 169 | + aclDestroyTensorList(outList); return 1); | ||
| 170 | + | ||
| 171 | + std::vector<float> outAHost(kRowsA * kCols, 0); | ||
| 172 | + std::vector<float> outBHost(kRowsB * kCols, 0); | ||
| 173 | + auto sizeA = kRowsA * kCols * sizeof(float); | ||
| 174 | + auto sizeB = kRowsB * kCols * sizeof(float); | ||
| 175 | + ret = aclrtMemcpy(outAHost.data(), sizeA, outADeviceAddr, sizeA, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 176 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: copy outA failed\n", caseName); return 1); | ||
| 177 | + ret = aclrtMemcpy(outBHost.data(), sizeB, outBDeviceAddr, sizeB, ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 178 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("[FAIL] %s: copy outB failed\n", caseName); return 1); | ||
| 179 | + | ||
| 180 | + auto expectedA = RepeatRows({1.0f, 2.0f, 3.0f}, kCols); | ||
| 181 | + auto expectedB = RepeatRows({4.0f, 5.0f, 6.0f, 7.0f}, kCols); | ||
| 182 | + int64_t errCntA = CompareBuf(outAHost, expectedA); | ||
| 183 | + int64_t errCntB = CompareBuf(outBHost, expectedB); | ||
| 184 | + | ||
| 185 | + int failures = 0; | ||
| 186 | + if (errCntA != 0 || errCntB != 0) { | ||
| 187 | + LOG_PRINT("[FAIL] %s: errCnt A=%ld, B=%ld\n", caseName, errCntA, errCntB); | ||
| 188 | + ++failures; | ||
| 189 | + } else { | ||
| 190 | + LOG_PRINT("[PASS] %s\n", caseName); | ||
| 191 | + } | ||
| 192 | + | ||
| 193 | + aclDestroyTensorList(outList); | ||
| 194 | + aclDestroyTensor(values); | ||
| 195 | + aclDestroyTensor(outA); | ||
| 196 | + aclDestroyTensor(outB); | ||
| 197 | + aclrtFree(inputDeviceAddr); | ||
| 198 | + aclrtFree(outADeviceAddr); | ||
| 199 | + aclrtFree(outBDeviceAddr); | ||
| 200 | + return failures; | ||
| 201 | +} | ||
| 202 | + | ||
| 203 | +int main() | ||
| 204 | +{ | ||
| 205 | + int32_t deviceId = 0; | ||
| 206 | + aclrtStream stream; | ||
| 207 | + auto ret = Init(deviceId, &stream); | ||
| 208 | + CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); | ||
| 209 | + | ||
| 210 | + int failures = 0; | ||
| 211 | + | ||
| 212 | + // ============================================================== | ||
| 213 | + // Case 1 & 2（对齐 torch: test_split_basic / test_split_with_prefix） | ||
| 214 | + // OFFSET_A = [0, 1, 2, 3], OFFSET_B = [0, 1, 3, 4] | ||
| 215 | + // 切分后 out_a (3, 8) = values_a: 每行全 1 / 2 / 3 | ||
| 216 | + // out_b (4, 8) = values_b: 每行全 4 / 5 / 6 / 7 | ||
| 217 | + // ============================================================== | ||
| 218 | + constexpr int64_t kCols = 8; | ||
| 219 | + constexpr int64_t kOffsetLen = 4; | ||
| 220 | + | ||
| 221 | + std::vector<int64_t> offsetsHost = {0, 1, 2, 3, 0, 1, 3, 4}; | ||
| 222 | + aclIntArray* offsetsArray = aclCreateIntArray(offsetsHost.data(), offsetsHost.size()); | ||
| 223 | + CHECK_RET(offsetsArray != nullptr, LOG_PRINT("aclCreateIntArray failed.\n"); return 1); | ||
| 224 | + | ||
| 225 | + // -------------------------------------------------------------- | ||
| 226 | + // Case 1: nPrefixToRight = 0（对齐 test_split_basic） | ||
| 227 | + // 输入 = concat(nPrefix=0) 产出: 1, 4, 2, 5, 6, 3, 7 | ||
| 228 | + // -------------------------------------------------------------- | ||
| 229 | + { | ||
| 230 | + std::vector<float> inputHost = RepeatRows({1.0f, 4.0f, 2.0f, 5.0f, 6.0f, 3.0f, 7.0f}, kCols); | ||
| 231 | + failures += RunSplitCase(inputHost, /*nPrefixToRight=*/0, offsetsArray, kOffsetLen, | ||
| 232 | + "split basic (nPrefix=0)", stream); | ||
| 233 | + } | ||
| 234 | + | ||
| 235 | + // -------------------------------------------------------------- | ||
| 236 | + // Case 2: nPrefixToRight = 1（对齐 test_split_with_prefix） | ||
| 237 | + // 输入 = concat(nPrefixFromRight=1) 产出: 4, 1, 5, 2, 6, 7, 3 | ||
| 238 | + // 切分后应还原成 values_a / values_b | ||
| 239 | + // -------------------------------------------------------------- | ||
| 240 | + { | ||
| 241 | + std::vector<float> inputHost = RepeatRows({4.0f, 1.0f, 5.0f, 2.0f, 6.0f, 7.0f, 3.0f}, kCols); | ||
| 242 | + failures += RunSplitCase(inputHost, /*nPrefixToRight=*/1, offsetsArray, kOffsetLen, | ||
| 243 | + "split with prefix (nPrefix=1)", stream); | ||
| 244 | + } | ||
| 245 | + | ||
| 246 | + aclDestroyIntArray(offsetsArray); | ||
| 247 | + | ||
| 248 | + aclrtDestroyStream(stream); | ||
| 249 | + aclrtResetDevice(deviceId); | ||
| 250 | + aclFinalize(); | ||
| 251 | + | ||
| 252 | + if (failures > 0) { | ||
| 253 | + LOG_PRINT("=== split_2d_jagged aclnn example FAILED, %d case(s) failed ===\n", failures); | ||
| 254 | + return 1; | ||
| 255 | + } | ||
| 256 | + LOG_PRINT("[aclnn] split_2d_jagged aclnn example PASSED\n"); | ||
| 257 | + return 0; | ||
| 258 | +} | ||
Aascendc/sequence/concat_jagged_tensor/split_2d_jagged/examples/torch/test_torch_split_2d_jagged.py+201-0
| @@ -0,0 +1,201 @@ | |||
| 1 | +#!/usr/bin/env python3 | ||
| 2 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | +# ============================================================================== | ||
| 16 | + | ||
| 17 | +""" | ||
| 18 | +split_2d_jagged 算子 PyTorch 接口基础用例 | ||
| 19 | + | ||
| 20 | +演示: | ||
| 21 | + 1. split_2d_jagged 基本前向切分 (nPrefixToRight=0) | ||
| 22 | + 2. nPrefixToRight 参数效果 (n_prefix=1) | ||
| 23 | + 3. split 前向 + backward 自动求导 | ||
| 24 | + 4. torch.compile / inductor 入图基础用例 | ||
| 25 | + | ||
| 26 | +用法: | ||
| 27 | + python3 test_split_2d_jagged.py | ||
| 28 | +""" | ||
| 29 | + | ||
| 30 | +import torch | ||
| 31 | +import ops_rec | ||
| 32 | + | ||
| 33 | +DEVICE = "npu:0" | ||
| 34 | +torch.npu.config.allow_internal_format = False | ||
| 35 | +torch.npu.set_device(DEVICE) | ||
| 36 | + | ||
| 37 | +# ============================================================================== | ||
| 38 | +# 公共数据 (与 forward example 同源，被 test_split_basic / test_split_with_prefix 复用) | ||
| 39 | +# ============================================================================== | ||
| 40 | + | ||
| 41 | +# nPrefixToRight=0 的 split 输入 = concat(values_a, values_b, nPrefix=0) 的产出 | ||
| 42 | +# 交错顺序: A[0], B[0], A[1], B[1:3], A[2], B[3] → 值 1, 4, 2, 5, 6, 3, 7 | ||
| 43 | +SPLIT_INPUT_PREFIX0 = torch.tensor([ | ||
| 44 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 45 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] | ||
| 46 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 47 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] | ||
| 48 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] | ||
| 49 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 50 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] | ||
| 51 | +], dtype=torch.float32) | ||
| 52 | + | ||
| 53 | +# nPrefixToRight=1 的 split 输入 = concat(values_a, values_b, nPrefix=1) 的产出 | ||
| 54 | +# B 每个 chunk 的前 1 行前移到该 chunk 最前: 4, 1, 5, 2, 6, 7, 3 | ||
| 55 | +SPLIT_INPUT_PREFIX1 = torch.tensor([ | ||
| 56 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] 第 1 行前移 | ||
| 57 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 58 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] 第 1 行前移 | ||
| 59 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 60 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] 剩余 | ||
| 61 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] 第 1 行前移 | ||
| 62 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 63 | +], dtype=torch.float32) | ||
| 64 | + | ||
| 65 | +OFFSET_A = torch.tensor([0, 1, 2, 3], dtype=torch.int64) | ||
| 66 | +OFFSET_B = torch.tensor([0, 1, 3, 4], dtype=torch.int64) | ||
| 67 | + | ||
| 68 | + | ||
| 69 | +# ============================================================================== | ||
| 70 | +# 预期输出（不依赖 golden 函数） | ||
| 71 | +# ============================================================================== | ||
| 72 | + | ||
| 73 | +# split 把输入按 chunk 边界交替切: | ||
| 74 | +# chunk0: sa=1(A), sb=1(B) → A[0]=[1], B[0]=[4] | ||
| 75 | +# chunk1: sa=1(A), sb=2(B) → A[1]=[2], B[1:3]=[5],[6] | ||
| 76 | +# chunk2: sa=1(A), sb=1(B) → A[2]=[3], B[3]=[7] | ||
| 77 | +# 两个用例的输入不同（prefix0 / prefix1），但切分结果都应恢复原始 values。 | ||
| 78 | +EXPECTED_OUT_A = torch.tensor([ | ||
| 79 | + [1, 1, 1, 1, 1, 1, 1, 1], # A[0] | ||
| 80 | + [2, 2, 2, 2, 2, 2, 2, 2], # A[1] | ||
| 81 | + [3, 3, 3, 3, 3, 3, 3, 3], # A[2] | ||
| 82 | +], dtype=torch.float32) | ||
| 83 | + | ||
| 84 | +EXPECTED_OUT_B = torch.tensor([ | ||
| 85 | + [4, 4, 4, 4, 4, 4, 4, 4], # B[0] | ||
| 86 | + [5, 5, 5, 5, 5, 5, 5, 5], # B[1] | ||
| 87 | + [6, 6, 6, 6, 6, 6, 6, 6], # B[2] | ||
| 88 | + [7, 7, 7, 7, 7, 7, 7, 7], # B[3] | ||
| 89 | +], dtype=torch.float32) | ||
| 90 | + | ||
| 91 | +# ============================================================================== | ||
| 92 | +# 测试用例 | ||
| 93 | +# ============================================================================== | ||
| 94 | + | ||
| 95 | +def test_split_basic(): | ||
| 96 | + """基本前向：nPrefixToRight=0，把 concat 产出切回 (out_a, out_b)""" | ||
| 97 | + result_a, result_b = ops_rec.ascendc.sequence.split_2d_jagged( | ||
| 98 | + SPLIT_INPUT_PREFIX0.to(DEVICE), 0, | ||
| 99 | + OFFSET_A.to(DEVICE), OFFSET_B.to(DEVICE), 0, 0 | ||
| 100 | + ) | ||
| 101 | + diff_a = (result_a.cpu() - EXPECTED_OUT_A).abs().max().item() | ||
| 102 | + diff_b = (result_b.cpu() - EXPECTED_OUT_B).abs().max().item() | ||
| 103 | + assert diff_a < 1e-6 and diff_b < 1e-6 | ||
| 104 | + print(f"[torch] split basic (nPrefix=0) PASSED: max diff A={diff_a:.6f}, B={diff_b:.6f}") | ||
| 105 | + | ||
| 106 | + | ||
| 107 | +def test_split_with_prefix(): | ||
| 108 | + """验证 nPrefixToRight = 1 的效果（nPrefixFromRight=1 的逆操作）""" | ||
| 109 | + result_a, result_b = ops_rec.ascendc.sequence.split_2d_jagged( | ||
| 110 | + SPLIT_INPUT_PREFIX1.to(DEVICE), 0, | ||
| 111 | + OFFSET_A.to(DEVICE), OFFSET_B.to(DEVICE), 0, 1 | ||
| 112 | + ) | ||
| 113 | + diff_a = (result_a.cpu() - EXPECTED_OUT_A).abs().max().item() | ||
| 114 | + diff_b = (result_b.cpu() - EXPECTED_OUT_B).abs().max().item() | ||
| 115 | + assert diff_a < 1e-6 and diff_b < 1e-6 | ||
| 116 | + print(f"[torch] split with prefix (nPrefix=1) PASSED: max diff A={diff_a:.6f}, B={diff_b:.6f}") | ||
| 117 | + | ||
| 118 | + | ||
| 119 | +def test_split_autograd(): | ||
| 120 | + """split_2d_jagged 前向 + backward 自动求导 | ||
| 121 | + | ||
| 122 | + 验证 split.backward 通过 autograd.py 注册的 concat_2d_jagged 反向路径 | ||
| 123 | + 正确把 (grad_a, grad_b) 拼回 values 的 grad。 | ||
| 124 | + """ | ||
| 125 | + torch.manual_seed(42) | ||
| 126 | + | ||
| 127 | + rows_a, rows_b, cols = 10, 15, 8 | ||
| 128 | + values_a = torch.randn(rows_a, cols, dtype=torch.float32, device=DEVICE) | ||
| 129 | + values_b = torch.randn(rows_b, cols, dtype=torch.float32, device=DEVICE) | ||
| 130 | + offset_a = torch.tensor([0, 3, 7, 10], dtype=torch.int64).npu() | ||
| 131 | + offset_b = torch.tensor([0, 4, 9, 15], dtype=torch.int64).npu() | ||
| 132 | + | ||
| 133 | + # 用 concat_2d_jagged 构造 split 的输入（nPrefix=0） | ||
| 134 | + with torch.no_grad(): | ||
| 135 | + concat_for_split = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 136 | + 0, values_a, values_b, offset_a, offset_b, False, 0 | ||
| 137 | + ) | ||
| 138 | + | ||
| 139 | + split_input = concat_for_split.requires_grad_(True) | ||
| 140 | + out_a, out_b = ops_rec.ascendc.sequence.split_2d_jagged( | ||
| 141 | + split_input, 0, offset_a, offset_b, 0, 0 | ||
| 142 | + ) | ||
| 143 | + (out_a.sum() + out_b.sum()).backward() | ||
| 144 | + | ||
| 145 | + # 因为 loss = (out_a + out_b).sum()，对输入每个元素的导数都是 1。 | ||
| 146 | + # split.backward 通过 concat_2d_jagged 把全 1 的 (grad_a, grad_b) 拼回去， | ||
| 147 | + # 所以 split_input.grad 应该全 1 且 shape 与输入一致。 | ||
| 148 | + grad = split_input.grad | ||
| 149 | + assert torch.allclose(grad, torch.ones_like(grad)), \ | ||
| 150 | + f"grad not all 1s: max diff from 1 = {(grad - 1).abs().max().item()}" | ||
| 151 | + | ||
| 152 | + print(f"[torch] split autograd: out_a.shape = {out_a.shape}, out_b.shape = {out_b.shape}") | ||
| 153 | + print(f"[torch] grad = 1 everywhere, shape = {grad.shape}") | ||
| 154 | + print("[torch] split autograd PASSED") | ||
| 155 | + | ||
| 156 | + | ||
| 157 | +def test_split_inductor(): | ||
| 158 | + """torch.compile / inductor 入图基础用例 | ||
| 159 | + | ||
| 160 | + 验证算子可被 dynamo 跟踪进图（dynamic=True, fullgraph=True）， | ||
| 161 | + 且 compiled 输出与 eager 一致。 | ||
| 162 | + """ | ||
| 163 | + torch.manual_seed(42) | ||
| 164 | + | ||
| 165 | + rows_a, rows_b, cols = 50, 80, 128 | ||
| 166 | + values_a = torch.randn(rows_a, cols, dtype=torch.float32, device=DEVICE) | ||
| 167 | + values_b = torch.randn(rows_b, cols, dtype=torch.float32, device=DEVICE) | ||
| 168 | + offset_a = torch.tensor([0, 15, 30, 50], dtype=torch.int64).npu() | ||
| 169 | + offset_b = torch.tensor([0, 20, 45, 80], dtype=torch.int64).npu() | ||
| 170 | + | ||
| 171 | + # 用 concat_2d_jagged 构造 split 的输入 | ||
| 172 | + with torch.no_grad(): | ||
| 173 | + concat_input = ops_rec.ascendc.sequence.concat_2d_jagged( | ||
| 174 | + 0, values_a, values_b, offset_a, offset_b, False, 0 | ||
| 175 | + ) | ||
| 176 | + | ||
| 177 | + eager_a, eager_b = ops_rec.ascendc.sequence.split_2d_jagged( | ||
| 178 | + concat_input, 0, offset_a, offset_b, 0, 0 | ||
| 179 | + ) | ||
| 180 | + | ||
| 181 | + compiled_fn = torch.compile( | ||
| 182 | + ops_rec.ascendc.sequence.split_2d_jagged, | ||
| 183 | + dynamic=True, fullgraph=True, | ||
| 184 | + ) | ||
| 185 | + compiled_a, compiled_b = compiled_fn(concat_input, 0, offset_a, offset_b, 0, 0) | ||
| 186 | + | ||
| 187 | + diff_a = (compiled_a.cpu() - eager_a.cpu()).abs().max().item() | ||
| 188 | + diff_b = (compiled_b.cpu() - eager_b.cpu()).abs().max().item() | ||
| 189 | + print(f"[torch] split inductor: eager A{eager_a.shape} B{eager_b.shape}, " | ||
| 190 | + f"compiled A{compiled_a.shape} B{compiled_b.shape}") | ||
| 191 | + print(f"[torch] split inductor: max diff A={diff_a:.6f}, B={diff_b:.6f}") | ||
| 192 | + assert diff_a < 1e-4 and diff_b < 1e-4 | ||
| 193 | + print("[torch] split inductor PASSED") | ||
| 194 | + | ||
| 195 | + | ||
| 196 | +if __name__ == "__main__": | ||
| 197 | + test_split_basic() | ||
| 198 | + test_split_with_prefix() | ||
| 199 | + test_split_autograd() | ||
| 200 | + test_split_inductor() | ||
| 201 | + print("\n=== All split_2d_jagged examples PASSED ===") | ||
| @@ -0,0 +1,78 @@ | |||
| 1 | +# ============================================================================ | ||
| 2 | +# concat_jagged_tensor_grad/op_host：算子 host 侧实现（OpDef + tiling + infershape） | ||
| 3 | +# ---------------------------------------------------------------------------- | ||
| 4 | +# 使用 CANN 的 npu_op_code_gen + npu_op_library 编译。 | ||
| 5 | +# npu_op_package 已在上层 CMakeLists.txt 调用。 | ||
| 6 | +set(MY_COMMON_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../../../../common) | ||
| 7 | + | ||
| 8 | +# OPS_REC_SOC_CXX_DEFS 由根 CMakeLists.txt 统一设置（SUPPORT_950 / SUPPORT_A2 / SUPPORT_A3） | ||
| 9 | + | ||
| 10 | +file(GLOB ops_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) | ||
| 11 | +file(GLOB ops_cc_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cc) | ||
| 12 | +list(APPEND ops_srcs ${ops_cc_srcs}) | ||
| 13 | + | ||
| 14 | +# 收集 SOC 架构特化源码（op_host/atlas950/ 或 op_host/atlasA2/ 目录下的额外实现） | ||
| 15 | +ops_rec_collect_arch_sources(ops_srcs) | ||
| 16 | + | ||
| 17 | +npu_op_code_gen( | ||
| 18 | + SRC ${ops_srcs} | ||
| 19 | + PACKAGE ${VENDOR_NAME} | ||
| 20 | + OUT_DIR ${ASCEND_AUTOGEN_PATH} | ||
| 21 | + COMPILE_OPTIONS | ||
| 22 | + ${OPS_REC_SOC_CXX_DEFS} | ||
| 23 | + -I${MY_COMMON_DIR}/log | ||
| 24 | + -I${MY_COMMON_DIR}/utils | ||
| 25 | +) | ||
| 26 | + | ||
| 27 | +file(GLOB autogen_aclnn_src ${ASCEND_AUTOGEN_PATH}/aclnn_*.cpp) | ||
| 28 | +set_source_files_properties(${autogen_aclnn_src} PROPERTIES GENERATED TRUE) | ||
| 29 | +npu_op_library(cust_opapi ACLNN | ||
| 30 | + ${autogen_aclnn_src} | ||
| 31 | +) | ||
| 32 | + | ||
| 33 | +target_include_directories(cust_opapi PRIVATE | ||
| 34 | + ${MY_COMMON_DIR}/log | ||
| 35 | + ${MY_COMMON_DIR}/utils | ||
| 36 | +) | ||
| 37 | + | ||
| 38 | +target_compile_options(cust_opapi PRIVATE | ||
| 39 | + -fvisibility=hidden | ||
| 40 | +) | ||
| 41 | + | ||
| 42 | +file(GLOB group_proto_src ${ASCEND_AUTOGEN_PATH}/group_op_proto/*.cc) | ||
| 43 | +file(GLOB proto_src ${ASCEND_AUTOGEN_PATH}/op_proto.cc) | ||
| 44 | +set_source_files_properties(${group_proto_src} PROPERTIES GENERATED TRUE) | ||
| 45 | +set_source_files_properties(${proto_src} PROPERTIES GENERATED TRUE) | ||
| 46 | +npu_op_library(cust_op_proto GRAPH | ||
| 47 | + ${ops_srcs} | ||
| 48 | + ${group_proto_src} | ||
| 49 | + ${proto_src} | ||
| 50 | +) | ||
| 51 | + | ||
| 52 | +target_include_directories(cust_op_proto PRIVATE | ||
| 53 | + ${MY_COMMON_DIR}/log | ||
| 54 | + ${MY_COMMON_DIR}/utils | ||
| 55 | +) | ||
| 56 | + | ||
| 57 | +target_compile_options(cust_op_proto PRIVATE | ||
| 58 | + -fvisibility=hidden | ||
| 59 | +) | ||
| 60 | + | ||
| 61 | +file(GLOB fallback_src ${ASCEND_AUTOGEN_PATH}/fallback_*.cpp) | ||
| 62 | +set_source_files_properties(${fallback_src} PROPERTIES GENERATED TRUE) | ||
| 63 | +npu_op_library(cust_optiling TILING | ||
| 64 | + ${ops_srcs} | ||
| 65 | + ${fallback_src} | ||
| 66 | +) | ||
| 67 | + | ||
| 68 | +target_include_directories(cust_optiling PRIVATE | ||
| 69 | + ${MY_COMMON_DIR}/log | ||
| 70 | + ${MY_COMMON_DIR}/utils | ||
| 71 | +) | ||
| 72 | + | ||
| 73 | +npu_op_package_add(${VENDOR_NAME} | ||
| 74 | + LIBRARY | ||
| 75 | + cust_optiling | ||
| 76 | + cust_opapi | ||
| 77 | + cust_op_proto | ||
| 78 | +) | ||
| @@ -0,0 +1,308 @@ | |||
| 1 | +/* Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); | ||
| 4 | +you may not use this file except in compliance with the License. | ||
| 5 | +You may obtain a copy of the License at | ||
| 6 | + | ||
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 | ||
| 8 | + | ||
| 9 | +Unless required by applicable law or agreed to in writing, software | ||
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, | ||
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 12 | +See the License for the specific language governing permissions and | ||
| 13 | + limitations under the License. | ||
| 14 | +==============================================================================*/ | ||
| 15 | + | ||
| 16 | +#include "split_jagged_tensor_tiling.h" | ||
| 17 | +#include "register/op_def_registry.h" | ||
| 18 | +#include "register/register.h" | ||
| 19 | +#include "tiling/platform/platform_ascendc.h" | ||
| 20 | +#include "ops_log.h" | ||
| 21 | + | ||
| 22 | +constexpr uint32_t RESERVED_UB_SIZE = 1024; | ||
| 23 | +constexpr uint32_t WORKSPACE_SIZE = 4 * 1024 * 1024; | ||
| 24 | +constexpr uint32_t ALIGN = 32; | ||
| 25 | +constexpr uint32_t MAX_TENSOR_SIZE = 24 * 1024; // 当Tensor长度大于此值时调整分核策略性能更优 | ||
| 26 | + | ||
| 27 | +namespace optiling { | ||
| 28 | + | ||
| 29 | +static ge::graphStatus TilingFunc(gert::TilingContext* context) | ||
| 30 | +{ | ||
| 31 | + // 获取硬件信息 | ||
| 32 | + uint64_t ubSize = 0; | ||
| 33 | + const auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); | ||
| 34 | + ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize); | ||
| 35 | + uint64_t maxUbSize = ubSize - RESERVED_UB_SIZE; | ||
| 36 | + uint64_t coreNum = ascendcPlatform.GetCoreNumAiv(); | ||
| 37 | + if (coreNum == 0) { | ||
| 38 | + OPS_LOG_E(context, "[ERROR] need more than 0 ai core"); | ||
| 39 | + return ge::GRAPH_FAILED; | ||
| 40 | + } | ||
| 41 | + | ||
| 42 | + // 计算偏移地址和大小 | ||
| 43 | + auto offsets = context->GetAttrs()->GetListInt(0); | ||
| 44 | + int32_t offsetLen = *context->GetAttrs()->GetInt(1); | ||
| 45 | + int32_t jtNum = *context->GetAttrs()->GetInt(2); | ||
| 46 | + int32_t nPrefixToRight = *context->GetAttrs()->GetInt(3); | ||
| 47 | + | ||
| 48 | + uint32_t sliceSize[MAX_SLICE_SIZE + 1] = {0}; // 当nPrefixFromRight > 0时,切片个数+1。 | ||
| 49 | + uint32_t inputOffsetBegin[MAX_SLICE_SIZE + 1] = {0}; | ||
| 50 | + uint32_t outputOffsetBegin[MAX_SLICE_SIZE + 1] = {0}; | ||
| 51 | + | ||
| 52 | + uint32_t indexSliceForRight = offsetLen - 1; | ||
| 53 | + int64_t inputOffset = 0; | ||
| 54 | + const int64_t* offsetsData = offsets->GetData(); | ||
| 55 | + | ||
| 56 | + uint32_t maxTensorLen = std::max(offsetsData[indexSliceForRight], offsetsData[2 * offsetLen - 1]); | ||
| 57 | + | ||
| 58 | + // 判断是否需要优化切片策略：Tensor长度较长且分片较小时 | ||
| 59 | + const bool needAdjustSliceStrategy = (offsetLen < coreNum) && (maxTensorLen > MAX_TENSOR_SIZE); | ||
| 60 | + | ||
| 61 | + if (nPrefixToRight == 0) { | ||
| 62 | + // [A] + [B] | ||
| 63 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 64 | + // A[:indexSliceForRight-1] | ||
| 65 | + outputOffsetBegin[i] = offsetsData[i]; | ||
| 66 | + sliceSize[i] = offsetsData[i + 1] - offsetsData[i]; | ||
| 67 | + // B[indexSliceForRight:] | ||
| 68 | + outputOffsetBegin[i + indexSliceForRight] = offsetsData[offsetLen + i]; | ||
| 69 | + sliceSize[i + indexSliceForRight] = offsetsData[offsetLen + i + 1] - offsetsData[offsetLen + i]; | ||
| 70 | + } | ||
| 71 | + } else { | ||
| 72 | + // [nPrefixFromB] + [A] + [remainB] | ||
| 73 | + // first offset [nPrefixFromB] | ||
| 74 | + int32_t bLength = offsetsData[offsetLen + 1] - offsetsData[offsetLen]; | ||
| 75 | + // nPrefixToRight | ||
| 76 | + sliceSize[indexSliceForRight] = std::min<int32_t>(nPrefixToRight, bLength); | ||
| 77 | + outputOffsetBegin[indexSliceForRight] = 0; | ||
| 78 | + | ||
| 79 | + // [A] + ([remainB] + [nPrefixFromB]) | ||
| 80 | + for (int i = 0; i < offsetLen - 2; i++) { | ||
| 81 | + // [A] | ||
| 82 | + outputOffsetBegin[i] = offsetsData[i]; | ||
| 83 | + sliceSize[i] = offsetsData[i + 1] - offsetsData[i]; | ||
| 84 | + // [remainB] | ||
| 85 | + int32_t remainB = std::max<int32_t>((bLength - nPrefixToRight), 0); | ||
| 86 | + // [next nPrefixFromB] | ||
| 87 | + bLength = offsetsData[offsetLen + i + 2] - offsetsData[offsetLen + i + 1]; | ||
| 88 | + int32_t prefixNextB = std::min<int32_t>(nPrefixToRight, bLength); | ||
| 89 | + | ||
| 90 | + outputOffsetBegin[i + offsetLen] = outputOffsetBegin[offsetLen + i - 1] + sliceSize[i + offsetLen - 1]; | ||
| 91 | + sliceSize[i + offsetLen] = remainB + prefixNextB; | ||
| 92 | + } | ||
| 93 | + // last offset [A] + [remainB] | ||
| 94 | + outputOffsetBegin[indexSliceForRight - 1] = offsetsData[indexSliceForRight - 1]; | ||
| 95 | + sliceSize[indexSliceForRight - 1] = offsetsData[indexSliceForRight] - offsetsData[indexSliceForRight - 1]; | ||
| 96 | + // B | ||
| 97 | + int32_t LastBLength = offsetsData[2 * offsetLen - 1] - offsetsData[2 * offsetLen - 2]; | ||
| 98 | + int32_t lastRemainB = std::max<int32_t>((LastBLength - nPrefixToRight), 0); | ||
| 99 | + | ||
| 100 | + outputOffsetBegin[2 * offsetLen - 2] = outputOffsetBegin[2 * offsetLen - 3] + sliceSize[2 * offsetLen - 3]; | ||
| 101 | + sliceSize[2 * offsetLen - 2] = lastRemainB; | ||
| 102 | + } | ||
| 103 | + | ||
| 104 | + uint64_t batchSize = (nPrefixToRight == 0) ? (jtNum * (offsetLen - 1)) : (jtNum * (offsetLen - 1) + 1); | ||
| 105 | + | ||
| 106 | + if (nPrefixToRight == 0) { | ||
| 107 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 108 | + inputOffsetBegin[i] = inputOffset; | ||
| 109 | + inputOffset += sliceSize[i]; | ||
| 110 | + inputOffsetBegin[i + offsetLen - 1] = inputOffset; | ||
| 111 | + inputOffset += sliceSize[i + offsetLen - 1]; | ||
| 112 | + } | ||
| 113 | + } else { | ||
| 114 | + inputOffsetBegin[offsetLen - 1] = inputOffset; | ||
| 115 | + inputOffset += sliceSize[offsetLen - 1]; | ||
| 116 | + for (int i = 0; i < offsetLen - 1; i++) { | ||
| 117 | + inputOffsetBegin[i] = inputOffset; | ||
| 118 | + inputOffset += sliceSize[i]; | ||
| 119 | + inputOffsetBegin[i + offsetLen] = inputOffset; | ||
| 120 | + inputOffset += sliceSize[i + offsetLen]; | ||
| 121 | + } | ||
| 122 | + } | ||
| 123 | + | ||
| 124 | + uint32_t oneSliceSize = maxTensorLen / coreNum; | ||
| 125 | + | ||
| 126 | + uint32_t resetIndexSliceForRight = indexSliceForRight; | ||
| 127 | + | ||
| 128 | + std::vector<uint32_t> readjustSliceSize; | ||
| 129 | + std::vector<uint32_t> readjustInputOffsetBegin; | ||
| 130 | + std::vector<uint32_t> readjustOutputOffsetBegin; | ||
| 131 | + | ||
| 132 | + if (needAdjustSliceStrategy) { | ||
| 133 | + for (int i = 0; i < batchSize; ++i) { | ||
| 134 | + if (i == indexSliceForRight) { | ||
| 135 | + resetIndexSliceForRight = readjustSliceSize.size(); | ||
| 136 | + } | ||
| 137 | + | ||
| 138 | + uint32_t inputSize = sliceSize[i]; | ||
| 139 | + uint32_t outputOffset = outputOffsetBegin[i]; | ||
| 140 | + uint32_t inputOffset = inputOffsetBegin[i]; | ||
| 141 | + | ||
| 142 | + if (inputSize <= oneSliceSize) { | ||
| 143 | + // 不需要切分，直接添加 | ||
| 144 | + readjustSliceSize.push_back(inputSize); | ||
| 145 | + readjustOutputOffsetBegin.push_back(outputOffset); | ||
| 146 | + readjustInputOffsetBegin.push_back(inputOffset); | ||
| 147 | + } else { | ||
| 148 | + // 需要切分成多个分片 | ||
| 149 | + uint32_t remaining = inputSize; | ||
| 150 | + uint32_t currentOutputOffset = outputOffset; | ||
| 151 | + uint32_t currentInputOffset = inputOffset; | ||
| 152 | + | ||
| 153 | + while (remaining > 0) { | ||
| 154 | + uint32_t chunkSize = std::min(remaining, oneSliceSize); | ||
| 155 | + readjustSliceSize.push_back(chunkSize); | ||
| 156 | + readjustOutputOffsetBegin.push_back(currentOutputOffset); | ||
| 157 | + readjustInputOffsetBegin.push_back(currentInputOffset); | ||
| 158 | + currentInputOffset += chunkSize; | ||
| 159 | + currentOutputOffset += chunkSize; | ||
| 160 | + remaining -= chunkSize; | ||
| 161 | + } | ||
| 162 | + } | ||
| 163 | + } | ||
| 164 | + } | ||
| 165 | + | ||
| 166 | + uint64_t resetBatchSize = batchSize; | ||
| 167 | + if (offsetLen < coreNum && maxTensorLen > MAX_TENSOR_SIZE) { | ||
| 168 | + resetBatchSize = readjustSliceSize.size(); | ||
| 169 | + } | ||
| 170 | + | ||
| 171 | + // 计算切分逻辑 | ||
| 172 | + uint64_t formerCore = 0; | ||
| 173 | + uint64_t tailCore = 0; | ||
| 174 | + uint64_t batchNumInTail = 0; | ||
| 175 | + uint64_t batchNumInFormer = 0; | ||
| 176 | + if (resetBatchSize < coreNum) { | ||
| 177 | + formerCore = resetBatchSize; | ||
| 178 | + batchNumInFormer = 1; | ||
| 179 | + } else { | ||
| 180 | + formerCore = resetBatchSize % coreNum; | ||
| 181 | + tailCore = coreNum - formerCore; | ||
| 182 | + batchNumInTail = resetBatchSize / coreNum; | ||
| 183 | + batchNumInFormer = batchNumInTail + 1; | ||
| 184 | + } | ||
| 185 | + | ||
| 186 | + // 计算copy最大行数 | ||
| 187 | + OPS_CHECK_PTR_NULL(context->GetRequiredInputShape(0), return ge::GRAPH_FAILED); | ||
| 188 | + OPS_LOG_E_IF_NULL("inputXDesc", context->GetInputDesc(0), return ge::GRAPH_FAILED); | ||
| 189 | + uint32_t inputColSize = context->GetRequiredInputShape(0)->GetOriginShape().GetDim(1); | ||
| 190 | + auto inputXDesc = context->GetInputDesc(0); | ||
| 191 | + ge::DataType xDtype{ge::DT_FLOAT}; | ||
| 192 | + xDtype = inputXDesc->GetDataType(); | ||
| 193 | + uint32_t ubMaxLength = maxUbSize / ge::GetSizeByDataType(xDtype); | ||
| 194 | + ubMaxLength = ubMaxLength & ~(ALIGN - 1); | ||
| 195 | + | ||
| 196 | + // 设置tiling | ||
| 197 | + SplitJaggedTensorTilingData tiling; | ||
| 198 | + tiling.set_jtNum(jtNum); | ||
| 199 | + tiling.set_inputColSize(inputColSize); | ||
| 200 | + tiling.set_ubMaxLength(ubMaxLength); | ||
| 201 | + tiling.set_formerCore(formerCore); | ||
| 202 | + tiling.set_tailCore(tailCore); | ||
| 203 | + tiling.set_batchNumInTail(batchNumInTail); | ||
| 204 | + tiling.set_batchNumInFormer(batchNumInFormer); | ||
| 205 | + | ||
| 206 | + if (needAdjustSliceStrategy) { | ||
| 207 | + // 填充剩余元素为 0 | ||
| 208 | + while (readjustSliceSize.size() < MAX_SLICE_SIZE + 1) { | ||
| 209 | + readjustSliceSize.push_back(0); | ||
| 210 | + readjustInputOffsetBegin.push_back(0); | ||
| 211 | + readjustOutputOffsetBegin.push_back(0); | ||
| 212 | + } | ||
| 213 | + tiling.set_indexSliceForRight(resetIndexSliceForRight); | ||
| 214 | + tiling.set_inputOffsetBegin(readjustInputOffsetBegin.data()); | ||
| 215 | + tiling.set_sliceSize(readjustSliceSize.data()); | ||
| 216 | + tiling.set_outputOffsetBegin(readjustOutputOffsetBegin.data()); | ||
| 217 | + | ||
| 218 | + } else { | ||
| 219 | + tiling.set_indexSliceForRight(indexSliceForRight); | ||
| 220 | + tiling.set_inputOffsetBegin(inputOffsetBegin); | ||
| 221 | + tiling.set_sliceSize(sliceSize); | ||
| 222 | + tiling.set_outputOffsetBegin(outputOffsetBegin); | ||
| 223 | + } | ||
| 224 | + | ||
| 225 | + context->SetBlockDim(coreNum); | ||
| 226 | + OPS_CHECK_PTR_NULL(context->GetRawTilingData(), return ge::GRAPH_FAILED); | ||
| 227 | + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity()); | ||
| 228 | + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); | ||
| 229 | + size_t* currentWorkspace = context->GetWorkspaceSizes(1); | ||
| 230 | + OPS_LOG_E_IF_NULL("currentWorkspace", currentWorkspace, return ge::GRAPH_FAILED); | ||
| 231 | + currentWorkspace[0] = WORKSPACE_SIZE; | ||
| 232 | + | ||
| 233 | + return ge::GRAPH_SUCCESS; | ||
| 234 | +} | ||
| 235 | +} // namespace optiling | ||
| 236 | + | ||
| 237 | +namespace ge { | ||
| 238 | +static ge::graphStatus InferShape(gert::InferShapeContext* context) | ||
| 239 | +{ | ||
| 240 | + auto offsets = context->GetAttrs()->GetListInt(0); | ||
| 241 | + int32_t offsetLen = *context->GetAttrs()->GetInt(1); | ||
| 242 | + int32_t jtNum = *context->GetAttrs()->GetInt(2); | ||
| 243 | + | ||
| 244 | + const gert::Shape* x_shape = context->GetInputShape(0); | ||
| 245 | + if (x_shape == nullptr) { | ||
| 246 | + OPS_LOG_E("[ERROR]", "InputShape should not be nullptr."); | ||
| 247 | + return ge::GRAPH_FAILED; | ||
| 248 | + } | ||
| 249 | + uint32_t dimNum = x_shape->GetDimNum(); | ||
| 250 | + uint32_t inputColSize = x_shape->GetDim(1); | ||
| 251 | + | ||
| 252 | + for (int i = 0; i < jtNum; i++) { | ||
| 253 | + gert::Shape* y_shape = context->GetOutputShape(i); | ||
| 254 | + if (y_shape == nullptr) { | ||
| 255 | + OPS_LOG_E("[ERROR]", "OutputShape should not be nullptr."); | ||
| 256 | + return ge::GRAPH_FAILED; | ||
| 257 | + } | ||
| 258 | + y_shape->SetDimNum(dimNum); | ||
| 259 | + y_shape->SetDim(0, offsets->GetData()[i * offsetLen]); | ||
| 260 | + y_shape->SetDim(1, inputColSize); | ||
| 261 | + } | ||
| 262 | + return ge::GRAPH_SUCCESS; | ||
| 263 | +} | ||
| 264 | +static ge::graphStatus InferDataType(gert::InferDataTypeContext* context) | ||
| 265 | +{ | ||
| 266 | + const auto inputDataType = context->GetInputDataType(0); | ||
| 267 | + int32_t jtNum = *context->GetAttrs()->GetInt(2); | ||
| 268 | + for (int i = 0; i < jtNum; i++) { | ||
| 269 | + context->SetOutputDataType(i, inputDataType); | ||
| 270 | + } | ||
| 271 | + return ge::GRAPH_SUCCESS; | ||
| 272 | +} | ||
| 273 | +} // namespace ge | ||
| 274 | + | ||
| 275 | +namespace ops { | ||
| 276 | +class SplitJaggedTensor : public OpDef { | ||
| 277 | +public: | ||
| 278 | + explicit SplitJaggedTensor(const char* name) : OpDef(name) | ||
| 279 | + { | ||
| 280 | + this->Input("values") | ||
| 281 | + .ParamType(REQUIRED) | ||
| 282 | + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32}) | ||
| 283 | + .FormatList({ge::FORMAT_ND}) | ||
| 284 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 285 | + this->Output("result") | ||
| 286 | + .ParamType(DYNAMIC) | ||
| 287 | + .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_BF16, ge::DT_INT32}) | ||
| 288 | + .FormatList({ge::FORMAT_ND}) | ||
| 289 | + .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 290 | + | ||
| 291 | + this->Attr("offsets").ListInt(); | ||
| 292 | + this->Attr("offsetLen").Int(0); | ||
| 293 | + this->Attr("jtNum").Int(0); | ||
| 294 | + this->Attr("nPrefixToRight").Int(0); | ||
| 295 | + | ||
| 296 | + this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); | ||
| 297 | + | ||
| 298 | + this->AICore().SetTiling(optiling::TilingFunc); | ||
| 299 | + this->AICore().AddConfig("ascend910b"); | ||
| 300 | + this->AICore().AddConfig("ascend910_93"); | ||
| 301 | +#ifdef SUPPORT_950 | ||
| 302 | + this->AICore().AddConfig("ascend950"); | ||
| 303 | +#endif | ||
| 304 | + } | ||
| 305 | +}; | ||
| 306 | + | ||
| 307 | +OP_ADD(SplitJaggedTensor); | ||
| 308 | +} // namespace ops | ||