vf.clear_spr#
产品支持情况#
Ascend 950PR/Ascend 950DT:支持
Atlas A3 训练系列产品/Atlas A3 推理系列产品:不支持
Atlas A2 训练系列产品/Atlas A2 推理系列产品:不支持
功能说明#
清除特殊寄存器(AddrReg寄存器),重置某些VF指令使用的累加寄存器。通常在 @pypto_pro.language.vector_function函数的开头或结尾调用。
函数原型#
clear_spr()
参数说明#
无参数。
约束说明#
无。
返回值说明#
无。
调用示例#
import os
import pypto_pro.language as pl
import torch
import torch_npu
@pl.vector_function
def example_vf(src_tile, dst_tile):
vf.clear_spr()
preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_FP32)
reg = vf.load_align(src_tile, 0)
vf.store_align(dst_tile, reg, preg)
@pl.jit()
def example_kernel(
a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32],
out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_FP32],
):
tf = pl.TileType(shape=[1, 64], dtype=pl.DT_FP32, target_memory=pl.MemorySpace.Vec)
in_a_grp = pl.make_tile_group(type=tf, addrs=0x0, mutex_ids=[0])
in_a = in_a_grp.current()
t_out_grp = pl.make_tile_group(type=tf, addrs=0x100, mutex_ids=[1])
t_out = t_out_grp.current()
with pl.section_vector():
pl.load(in_a, a, [0, 0])
example_vf(in_a, t_out)
pl.store(out, t_out, [0, 0])
def test_example():
device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0))
device = f"npu:{device_id}"
core_nums = 1
torch.npu.set_device(device)
a = torch.randn([1, 64], device=device, dtype=torch.float32)
out = torch.empty([1, 64], device=device, dtype=torch.float32)
example_kernel[None, core_nums](a, out)
torch.npu.synchronize()
torch.testing.assert_close(out, a, rtol=1e-5, atol=1e-5)
if __name__ == "__main__":
test_example()
print("PASSED")