triton.language.view
1 功能作用说明
创建张量的视图,改变形状但不复制数据,类似于reshape,但更强调视图的概念,保持数据在内存中的连续性。
语法:
triton.language.view(input, shape)- 函数调用形式input.view(shape)- 成员函数形式
功能:
创建张量的视图,改变形状但不复制数据
类似于reshape,但更强调视图的概念
保持数据在内存中的连续性
2 参数规格
2.1 参数说明
参数名 |
类型 |
必需 |
说明 |
|---|---|---|---|
input |
tensor |
是 |
输入张量 |
shape |
List[int] |
是 |
目标形状 |
返回值:
类型: tensor
形状: 与shape参数指定的目标形状相同
数据类型: 与输入张量相同
内存布局: 与输入张量在内存中连续
约束条件:
输入和输出张量的总元素数必须相等
输出张量必须与输入张量在内存中连续
2.2 DataType支持表
支持情况 |
int8 |
int16 |
int32 |
int64 |
uint8 |
uint16 |
uint32 |
uint64 |
float16 |
float32 |
bfloat16 |
float8e4 |
float8e5 |
float64 |
bool |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
Ascend A2/A3 |
✓ |
✓ |
✓ |
✓ |
✓ |
× |
× |
× |
✓ |
✓ |
✓ |
× |
× |
× |
✓ |
GPU支持 |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
✓ |
2.3 Shape支持表
支持任意维度数、任意形状大小。
2.4 特殊限制说明
无
2.5 使用方法
import torch
import triton
import triton.language as tl
@triton.jit
def view_example(out_ptr):
# 创建2x3x4的张量
x = tl.zeros([2, 3, 4], dtype=tl.float32)
# 创建视图,变成6x4
y = tl.view(x, [6, 4])
# 将结果写回外部张量
offs = tl.arange(0, 6)[:, None] * 4 + tl.arange(0, 4)[None, :]
tl.store(out_ptr + offs, y)
## 调用示例
out = torch.empty((6, 4), dtype=torch.float32, device="npu")
view_example[(1,)](out)
print(out.shape) # 输出: torch.Size([6, 4])