支持tensor_descriptor算子（make/load/store_tensor_descriptor）

## 一、需求场景&价值
`tensor_descriptor` 系列算子用于在 Triton 中显式描述张量的**布局信息**（shape、stride、block 结构等），并支持将该描述符用于后续的内存读写操作（load/store）等等。

**价值**：可以提前构建张量布局信息，避免在每次内存操作时重复计算地址，提高指令效率。

## 二、算子设计
### 1. make_tensor_descriptor
简介：创建张量描述符对象
表达式：
```python
triton.language.make_tensor_descriptor(base: tensor, shape: List[tensor], strides: List[tensor], block_shape: List[constexpr]) -> tensor_descriptor
```
参数说明：

| 参数名           | 类型                | 说明                                        |
| ------------- | ----------------- | ----------------------------------------- |
| `base`        | `tensor`          | 张量数据的起始地址                                 |
| `shape`       | `List[tensor]`    | 张量各维度大小                                   |
| `strides`     | `List[tensor]`    | 各维度的 stride，前面的维度必须是16字节的整数倍、最后一维必须是连续存储的 |
| `block_shape` | `List[constexpr]` | 从全局内存加载/存储的块的形状                           |

返回值：tensor_descriptor描述符对象（不可以直接算数运算，需配合load/store使用）

### 2. load_tensor_descriptor
简介：从张量描述符中读取张量数据
表达式：
```python
triton.language.load_tensor_descriptor(desc: tensor_descriptor_base, offsets: Sequence[constexpr | tensor])→ tensor
```
参数说明：

| 参数名       | 类型                              | 说明    |
| --------- | ------------------------------- | ----- |
| `desc`    | `tensor_descriptor_base`        | 张量描述符 |
| `offsets` | `Sequence[constexpr \| tensor]` | 元素偏移  |

返回值：tensor

### 3. store_tensor_descriptor
简介：向张量描述符指向的张量位置写入数据
表达式：
```python
triton.language.store_tensor_descriptor(desc: tensor_descriptor_base, offsets: Sequence[constexpr | tensor], value: tensor) → tensor
```
参数说明：

| 参数名       | 类型                              | 说明    |
| --------- | ------------------------------- | ----- |
| `desc`    | `tensor_descriptor_base`        | 张量描述符 |
| `offsets` | `Sequence[constexpr \| tensor]` | 元素偏移  |
| `value`   | `tensor`                        | 写入数据  |

返回值：tensor

参考文档：
https://triton-lang.org/main/python-api/generated/triton.language.make_tensor_descriptor.html

## 三、改进建议
### 目标
1. 为 Triton 增加 `make/load/store_tensor_descriptor` 三个 IR 操作及其 Python API
2. 确保新增功能不影响现有 Triton 算子与优化
### 方案对比
#### 方案一：基于 triton-patch 扩展
1. Triton IR 与方言扩展
* 核心定义IR文件：`triton_patch/include/triton/Dialect/Triton/IR/TritonOps.td`、`triton_patch/lib/Dialect/Triton/IR/Ops.cpp`
* python C++绑定：`triton_patch/python/src/ir.cc` 
* python 层接口注册和实现：
	注册用户API - `triton_patch/python/triton_patch/language/core.py`
	实现语义函数 - `triton_patch/python/triton_patch/language/semantic.py`

2. TritonTolinalg 转换逻辑（修改triton-adapter）
- `MakeTensorDescriptorOp`：仅作为元信息容器，无需直接 Lower
- `DescriptorLoadOp` / `DescriptorStoreOp`：Lower 为 `make_tensor_ptr + load/store` 的组合

✅ 优点：
* 改动范围小：只需在 patch 和 adapter 中新增对应 IR 与转换逻辑，不影响整体编译和现有算子
* 风险可控：不涉及底层架构和依赖大幅调整，短期内更容易验证和交付
* 快速落地：适合当前阶段尽快接入 tensor_descriptor 功能

⚠️ 缺点：
* 依赖 triton-patch：与上游社区存在差异，后续需要持续维护 patch
* 功能受限：部分算子逻辑依赖 Triton 新版本，旧版本 patch 难以完整支持
* 潜在兼容性问题：随着社区演进，维护 patch 成本可能逐步增加

#### 方案二：基于新版本 Triton
1. 升级 third_party 中的 Triton 至 `3.4.0`
2. 在新版本基础上增加 `TritonTolinalg` 转换逻辑

✅ 优点：
* 对齐社区主线：减少与上游差异，后续维护与升级更轻量
* 长期可持续：便于直接继承社区新特性、优化 pass 和后端适配
* 扩展性强：新版本语义和 IR 更规范，有利于未来在 tensor_descriptor 之外拓展更多功能

⚠️ 缺点：
* 编译流程变更：升级 Triton 至 3.4 需要调整现有的编译脚本、依赖管理方式
* 语义层（semantic）重构：新版本 Triton 的 semantic 定义方式有明显变化，需要重新对接、适配
* 存量算子逻辑变动：部分已有算子的行为或实现细节在新版本中有所调整（如cast等），需要逐一回归验证
* C++依赖文件变化：部分接口/文件已被移除或重构（如 TritonTypeInterfaces.td），需重新对接 Dialect include 与 lib
* 短期成本高：升级涉及面广，不仅需要升级，还要重写和适配大量已有逻辑，短期内会显著增加研发和测试成本

### 建议
- **短期**：采用方案一，快速接入 tensor_descriptor
- **中长期**：跟进 Triton 主线版本（3.4.0），逐步迁移到方案二

Ascend/triton-ascend
暂停

内容风险标识

评论 (0)

Ascend/triton-ascend暂停 .gitee-modal { width: 500px !important; }

内容风险标识