decomp pow for LVP (#14084)

test failed due to undefined behavior, so use decomp instead
This commit is contained in:
chenyu
2026-01-09 10:50:28 -05:00
committed by GitHub
parent 0fbc551622
commit 74fa3c7d09
2 changed files with 11 additions and 5 deletions

View File

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
import numpy as np
from typing import List, Callable
import torch
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM, EMULATE
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM, EMULATE
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.device import is_dtype_supported
@@ -235,8 +235,7 @@ class TestOps(unittest.TestCase):
def test_unfold(self):
helper_test_op([(8,)], lambda x: x.unfold(0, 2, 1))
helper_test_op([(8,)], lambda x: x.unfold(0, 2, 2))
# TODO: something is wrong with unfold
if not getenv("TINY_BACKEND"): helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
helper_test_op([(8,)], lambda x: x.unfold(0, 7, 3))
helper_test_op([(3,3,3)], lambda x: x.unfold(2, 2, 8))
helper_test_op([(3,3,3)], lambda x: x.unfold(1, 0, 8))
helper_test_op([(3,3,3,3,3)], lambda x: x.unfold(-1, 2, 2))
@@ -701,10 +700,14 @@ class TestOps(unittest.TestCase):
def test_pow_zero_tensor(self):
helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]])
# TODO: fix WEBGPU and LVP
if Device.DEFAULT != "WEBGPU" and not CPU_LVP:
# TODO: fix WEBGPU
if Device.DEFAULT != "WEBGPU":
helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.3]])
helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [-0.3]])
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU issue")
def test_exp2_log2_zero_times_negative(self):
# gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
def test_pow_zero_const(self):
helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])

View File

@@ -234,6 +234,9 @@ class LVPRenderer(NIRRenderer):
has_shared = False
global_max = (1, 0, 0)
nir_options = mesa.lvp_nir_options
# gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
# https://gitlab.freedesktop.org/mesa/mesa/-/blob/c200b18e876468b51fe80d9660f612dc03a5138e/src/gallium/auxiliary/gallivm/lp_bld_arit.c#L2972
code_for_op = {k:v for k,v in NIRRenderer.code_for_op.items() if k != Ops.EXP2}
param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz: