triton-lang/triton
[Question] How to force certain computations to occur in float16?
开放
#1,090 创建于 2023年1月24日
help wanted
仓库指标
- 星标
- (19,995 个星标)
- PR 合并指标
- (平均合并 2天 18小时) (30 天内合并 185 个 PR)
描述
I have the following kernel:
@triton.jit
def unmasked_blend_kernel(
img1_ptr, img2_ptr, ratio: float, total_items: int, BLOCK_SIZE: tl.constexpr
):
block_idx = tl.program_id(0)
offset = block_idx * BLOCK_SIZE
item_element_idxs = tl.arange(0, BLOCK_SIZE)
img1_item_ptrs = img1_ptr + offset + item_element_idxs
img2_item_ptrs = img2_ptr + offset + item_element_idxs
if offset + BLOCK_SIZE > total_items:
mask = item_element_idxs < total_items % BLOCK_SIZE
img1 = tl.load(img1_item_ptrs, mask=mask, other=0)
img2 = tl.load(img2_item_ptrs, mask=mask, other=0)
out = ratio * img1 + (1.0 - ratio) * img2
tl.store(img1_item_ptrs, out, mask=mask)
else:
img1 = tl.load(img1_item_ptrs)
img2 = tl.load(img2_item_ptrs)
out = ratio * img1 + (1.0 - ratio) * img2
tl.store(img1_item_ptrs, out)
Given that this kernel is blending 2 uint8 tensors, I figured it would make sense to do the computation out = ratio * img1 + (1.0 - ratio) * img2 since I don't need high precision.
However, I can't figure out how to force the Triton compiler to do the computations in FP16. I tried doing:
img1 = tl.load(img1_item_ptrs, mask=mask, other=0).to(tl.float16)
img2 = tl.load(img2_item_ptrs, mask=mask, other=0).to(tl.float16)
but from looking at the generated PTX, it just seems like the float16s are just converted to float32s before the multiplication occurs.
Is there a way I can force the Triton compiler to make the multiplications float16? (Or is this impossible for a good reason; i.e maybe this is not what I actually want)?