forked from tinygrad/tinygrad
Merge branch 'master' into move_gates_to_load_store
This commit is contained in:
George Hotz
@@ -44,7 +44,7 @@ jobs:
|
||||
python3 -c "from tinygrad.runtime.autogen import opencl"
|
||||
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
|
||||
python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
|
||||
python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2, fw, navi_offsets, vega_offsets, regs"
|
||||
python3 -c "from tinygrad.runtime.autogen.am import *"
|
||||
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
|
||||
python3 -c "from tinygrad.runtime.autogen import llvm"
|
||||
python3 -c "from tinygrad.runtime.autogen import webgpu"
|
||||
|
||||
@@ -417,7 +417,7 @@ jobs:
|
||||
llvm: 'true'
|
||||
- name: Test openpilot model kernel count and gate usage
|
||||
run: |
|
||||
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
|
||||
ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
|
||||
- name: Test openpilot CL compile fp16
|
||||
run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
|
||||
- name: Test openpilot CL compile fp32 (test correctness)
|
||||
|
||||
@@ -1419,7 +1419,10 @@ def train_llama3():
|
||||
|
||||
for p in optim.params:
|
||||
grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
|
||||
p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
|
||||
if isinstance(p.device, tuple) and p.uop.axis is not None:
|
||||
p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device[0]).shard_(p.device, axis=p.uop.axis).contiguous()
|
||||
else:
|
||||
p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
|
||||
grads = [p.grad for p in optim.params]
|
||||
|
||||
scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
|
||||
|
||||
@@ -81,7 +81,7 @@ class GradAccClipAdamW(Optimizer):
|
||||
if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
|
||||
if t.dtype in dtypes.fp8s:
|
||||
from examples.mlperf.models.flat_llama import FP8_MAX
|
||||
amax = new_w.float().abs().flatten(1).max(1).detach() # per-layer amax for (n_layers, out, in)
|
||||
amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach() # per-layer amax for (n_layers, out, in)
|
||||
scale = FP8_MAX / (amax + 1e-8)
|
||||
fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
|
||||
if hasattr(t, '_inv_scale'):
|
||||
|
||||
+12
-5
@@ -2,7 +2,6 @@
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
@@ -10,14 +9,22 @@ export DEVICE_IN_FUNCTION_BUG=1
|
||||
export DEBUG=${DEBUG:-2}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-1}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
|
||||
export FP8=${FP8:-1}
|
||||
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
|
||||
export FAST_CE=${FAST_CE:-0}
|
||||
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
|
||||
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
|
||||
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
|
||||
export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
export DP=${DP:-1} MP=${MP:-8}
|
||||
export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
|
||||
export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
|
||||
export GBS=$((BS * GRADIENT_ACC_STEPS))
|
||||
|
||||
export MODEL="llama3"
|
||||
export BASEDIR="/raid/datasets/c4/"
|
||||
@@ -30,7 +37,7 @@ export DATA_SEED=${DATA_SEED:-5760}
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
export FAKEDATA=1 BENCHMARK=10
|
||||
export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=2
|
||||
fi
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ export WQKV=${WQKV:-1}
|
||||
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
|
||||
export FP8=${FP8:-1}
|
||||
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
|
||||
export FAST_CE=${FASE_CE:-1}
|
||||
export FAST_CE=${FAST_CE:-1}
|
||||
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
|
||||
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
|
||||
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
|
||||
|
||||
+10
-3
@@ -2,7 +2,6 @@
|
||||
|
||||
export PYTHONPATH="."
|
||||
export DEV=${DEV:-AMD}
|
||||
export EMULATE="AMD_CDNA4"
|
||||
export CHECK_OOB=0
|
||||
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
|
||||
export DEVICE_IN_FUNCTION_BUG=1
|
||||
@@ -10,9 +9,17 @@ export DEVICE_IN_FUNCTION_BUG=1
|
||||
export DEBUG=${DEBUG:-2}
|
||||
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
|
||||
export ALL2ALL=${ALL2ALL:-1}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-0}
|
||||
export USE_ATOMICS=${USE_ATOMICS:-1}
|
||||
export ASM_GEMM=${ASM_GEMM:-1}
|
||||
export WQKV=${WQKV:-1}
|
||||
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
|
||||
export FP8=${FP8:-1}
|
||||
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
|
||||
export FAST_CE=${FAST_CE:-0}
|
||||
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
|
||||
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
|
||||
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
|
||||
export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
|
||||
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
|
||||
|
||||
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
|
||||
@@ -35,7 +42,7 @@ export DATA_SEED=${DATA_SEED:-5760}
|
||||
export JITBEAM=${JITBEAM:-3}
|
||||
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
|
||||
|
||||
export FAKEDATA=1 BENCHMARK=10
|
||||
export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
|
||||
if [ -z "$FULL_LAYERS" ]; then
|
||||
export LLAMA_LAYERS=2
|
||||
fi
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ export WQKV=${WQKV:-1}
|
||||
export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
|
||||
export FP8=${FP8:-1}
|
||||
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
|
||||
export FAST_CE=${FASE_CE:-1}
|
||||
export FAST_CE=${FAST_CE:-1}
|
||||
export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
|
||||
export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
|
||||
export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
|
||||
|
||||
@@ -34,13 +34,12 @@ def dname_of(device) -> str:
|
||||
return device.split(":")[0] if isinstance(device, str) else device
|
||||
|
||||
def alloc_like(shape, dtype, device, axis=None) -> Tensor:
|
||||
if isinstance(device, tuple):
|
||||
if axis is None: return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
|
||||
if isinstance(device, tuple) and axis is not None:
|
||||
return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
|
||||
return Tensor.invalids(*shape, dtype=dtype, device=device)
|
||||
|
||||
def alloc_local(shape, dtype, device) -> Tensor:
|
||||
if isinstance(device, tuple):
|
||||
def alloc_local(shape, dtype, device, axis=None) -> Tensor:
|
||||
if isinstance(device, tuple) and axis is not None:
|
||||
return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
|
||||
return Tensor.invalids(*shape, dtype=dtype, device=device)
|
||||
|
||||
|
||||
@@ -41,10 +41,9 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
|
||||
_, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
|
||||
device = xw13.device
|
||||
axis = xw13.axis if isinstance(device, tuple) else None
|
||||
if isinstance(device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
|
||||
grad_xw13 = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
|
||||
grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3, device, axis)
|
||||
grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32, device)
|
||||
grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32, device, axis)
|
||||
grad_amax_state_t = Tensor(grad_amax_state, device=device)
|
||||
fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
|
||||
grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
|
||||
@@ -66,9 +65,8 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
|
||||
assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
|
||||
HIDDEN = H2 // 2
|
||||
axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
|
||||
if isinstance(xw13.device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
|
||||
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, xw13.device, axis)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, xw13.device)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, xw13.device, axis)
|
||||
fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
|
||||
fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state,
|
||||
fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
|
||||
|
||||
@@ -63,7 +63,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st
|
||||
MBS, SEQ, HIDDEN = x_normed_u.shape
|
||||
axis = x_normed_u.axis if isinstance(device, tuple) else None
|
||||
grad_x = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, device, axis)
|
||||
grad_weight_partial = alloc_local((NUM_WG, HIDDEN), dtypes.float32, device)
|
||||
grad_weight_partial = alloc_local((NUM_WG, HIDDEN), dtypes.float32, device, axis)
|
||||
grad_h_from_fp8 = None
|
||||
grad_weight_uop = None
|
||||
if fp8_grad_u is not None:
|
||||
@@ -119,11 +119,11 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
|
||||
assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
|
||||
MBS, SEQ, HIDDEN = x.shape
|
||||
axis = x.uop.axis if isinstance(x.device, tuple) else None
|
||||
if isinstance(x.device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
|
||||
if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
|
||||
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis)
|
||||
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
|
||||
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
|
||||
fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
|
||||
fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
|
||||
fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
|
||||
@@ -139,12 +139,12 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
|
||||
assert x.shape == residual.shape
|
||||
MBS, SEQ, HIDDEN = x.shape
|
||||
axis = x.uop.axis if isinstance(x.device, tuple) else None
|
||||
if isinstance(x.device, tuple): assert axis in (0, 1), f"unsupported sharding axis={axis}"
|
||||
if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
|
||||
fp8_out = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype, x.device, axis)
|
||||
h_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
|
||||
x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
|
||||
rrms_out = alloc_like((MBS, SEQ), dtypes.float32, x.device, axis)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device)
|
||||
amax_buf = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
|
||||
fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
|
||||
fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
|
||||
fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
|
||||
|
||||
@@ -49,7 +49,7 @@ def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3)
|
||||
assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
|
||||
axis = x.uop.axis if isinstance(x.device, tuple) else None
|
||||
fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
|
||||
amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device)
|
||||
amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
|
||||
fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device))
|
||||
fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
|
||||
fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
|
||||
|
||||
+1
-2
@@ -251,8 +251,7 @@ select = [
|
||||
"F541",
|
||||
"F841",
|
||||
]
|
||||
"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E722", "E731", "F821", "A006", "A002", "F811"]
|
||||
"tinygrad/runtime/autogen/amd/**/*.py" = ["E501"]
|
||||
"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E731", "F821", "A006", "A002", "F811", "F822"]
|
||||
"test/amd/**/*.py" = ["F403", "F405"]
|
||||
|
||||
[tool.ruff.format]
|
||||
|
||||
Vendored
+5
-7
@@ -7,7 +7,7 @@ import z3
|
||||
from tinygrad import Variable, dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.uop.validate import uops_to_z3
|
||||
from tinygrad.helpers import DEBUG, Context
|
||||
from tinygrad.helpers import DEBUG
|
||||
|
||||
seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
|
||||
print(f"Seed: {seed}", flush=True)
|
||||
@@ -56,8 +56,7 @@ if __name__ == "__main__":
|
||||
v = [u1,u2,u3]
|
||||
expr = random_int_expr(6)
|
||||
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
simplified_expr = expr.simplify()
|
||||
simplified_expr = expr.simplify()
|
||||
|
||||
solver = z3.Solver(ctx=z3.Context())
|
||||
solver.set(timeout=5000) # some expressions take very long verify, but its very unlikely they actually return sat
|
||||
@@ -74,10 +73,9 @@ if __name__ == "__main__":
|
||||
m = solver.model()
|
||||
n1, n2, n3 = m[v1], m[v2], m[v3]
|
||||
u1_val, u2_val, u3_val = u1.const_like(n1.as_long()), u2.const_like(n2.as_long()), u3.const_like(n3.as_long())
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
|
||||
rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
|
||||
if num==rn: print("z3 found a mismatch but the expressions are equal!!")
|
||||
num = expr.simplify().substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
|
||||
rn = expr.substitute({u1:u1_val, u2:u2_val, u3:u3_val}).ssimplify()
|
||||
if num==rn: print("z3 found a mismatch but the expressions are equal!!")
|
||||
assert False, f"mismatched {expr.render()} at v1={m[v1]}; v2={m[v2]}; v3={m[v3]} = {num} != {rn}\n" +\
|
||||
"Reproduce with:\n" +\
|
||||
f"v1=Variable(\"{u1.arg[0]}\", {u1.arg[1]}, {u1.arg[2]})\n" +\
|
||||
|
||||
+2
-3
@@ -2,7 +2,7 @@ import random, sys
|
||||
import z3
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.uop.validate import uops_to_z3
|
||||
from tinygrad.helpers import DEBUG, Context, colored
|
||||
from tinygrad.helpers import DEBUG, colored
|
||||
|
||||
seed = int(sys.argv[1]) if len(sys.argv) > 1 else random.randint(0, 100)
|
||||
print(f"Seed: {seed}", flush=True)
|
||||
@@ -36,8 +36,7 @@ if __name__ == "__main__":
|
||||
variable_names += [f"r{i}" for i in range(num_ranges)]
|
||||
expr = get_random_expr(ranges, factors)
|
||||
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
simplified_expr = expr.simplify()
|
||||
simplified_expr = expr.simplify()
|
||||
|
||||
if DEBUG>=1:
|
||||
print(expr.render(simplify=False), " --> ", simplified_expr.render(simplify=False))
|
||||
|
||||
@@ -1,12 +1,18 @@
|
||||
import unittest, itertools
|
||||
|
||||
from tinygrad.codegen.late.devectorizer import load_store_indexing
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
from tinygrad.uop.symbolic import simplify_valid
|
||||
from tinygrad.uop.ops import UOp, Ops, graph_rewrite
|
||||
from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
|
||||
from tinygrad.helpers import Context
|
||||
from test.helpers import full_rewrite
|
||||
from test.null.test_uop_symbolic import check_uop_against_string
|
||||
|
||||
# symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
|
||||
def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
|
||||
# image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
|
||||
def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+load_store_indexing, name="simplify_image_idx")
|
||||
|
||||
def get_gated_load_uop(valid:UOp, idx:UOp):
|
||||
return UOp(Ops.LOAD, dtypes.float, (
|
||||
UOp(Ops.PARAM, dtypes.float.ptr(), arg=0).index(idx.valid(valid), ptr=True),
|
||||
@@ -47,11 +53,10 @@ class TestHelpers(unittest.TestCase):
|
||||
|
||||
class TestValidIdxSimplification(unittest.TestCase):
|
||||
def check(self, load, sidx, svalid, extra=()):
|
||||
with Context(NOOPT=1, SPEC=0):
|
||||
load = full_rewrite(UOp.sink(load, *extra)).src[0]
|
||||
idx, valid = load.src[0].src[1], load.src[0].src[2]
|
||||
check_uop_against_string(self, idx, sidx)
|
||||
check_uop_against_string(self, valid, svalid)
|
||||
load = simplify_valid_idx(UOp.sink(load, *extra)).src[0]
|
||||
off = load.src[0].src[1]
|
||||
check_uop_against_string(self, off.get_idx(), sidx)
|
||||
check_uop_against_string(self, off.get_valid(), svalid)
|
||||
|
||||
def test_cumsum(self):
|
||||
gidx0 = Special("gidx0", 5)
|
||||
@@ -216,18 +221,18 @@ class TestValidIdxSimplification(unittest.TestCase):
|
||||
|
||||
class TestImageSimplification(unittest.TestCase):
|
||||
def check(self, load, svalid, sidx0, sidx1):
|
||||
with Context(NOOPT=1, SPEC=0):
|
||||
load = full_rewrite(load.sink()).src[0]
|
||||
idx = load.src[0].src[1]
|
||||
load = simplify_image_idx(load.sink()).src[0]
|
||||
off = load.src[0].src[1]
|
||||
idx = off.get_idx()
|
||||
self.assertEqual(idx.op, Ops.STACK)
|
||||
self.assertEqual(len(idx.src), 2)
|
||||
idx0, idx1 = idx.src[0], idx.src[1]
|
||||
check_uop_against_string(self, idx0, sidx0)
|
||||
check_uop_against_string(self, idx1, sidx1)
|
||||
if svalid is not None:
|
||||
check_uop_against_string(self, load.src[0].src[2], svalid)
|
||||
check_uop_against_string(self, off.get_valid(), svalid)
|
||||
else:
|
||||
self.assertEqual(len(load.src[0].src), 2, "svalid is None but load still has a valid")
|
||||
self.assertEqual(off.get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")
|
||||
|
||||
def test_idx_gt_c(self):
|
||||
# (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
|
||||
@@ -447,12 +452,12 @@ class TestImageSimplification(unittest.TestCase):
|
||||
load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
|
||||
self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
|
||||
|
||||
# TODO: this is the same idx as above, but simplifying idx too early makes it hard to drop the valid
|
||||
# same idx, written without the inline simplification of the inner div/mod
|
||||
alu0 = ((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)
|
||||
alu1 = (lidx0*2+r0+-3)
|
||||
valid = ((lidx1<7)&((((lidx0*2+r0)<3)!=1)&((lidx0*2+r0)<35)))
|
||||
load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
|
||||
self.check(load, "(lidx1<7)", "((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)", "(lidx0*2+r0+-3)")
|
||||
self.check(load, None, "(lidx1*128+gidx0//2+144)", "(lidx0*2+r0+-3)")
|
||||
|
||||
def test_simplify8(self):
|
||||
# from openpilot compile3, kernel r_4_16_8_16_4_4_3_3n1
|
||||
|
||||
@@ -1,16 +1,8 @@
|
||||
import unittest
|
||||
from tinygrad import Variable
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
|
||||
class TestFuzzFailure(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.context = Context(CORRECT_DIVMOD_FOLDING=1)
|
||||
self.context.__enter__()
|
||||
|
||||
def tearDown(self):
|
||||
self.context.__exit__(None, None, None)
|
||||
|
||||
def test_fuzz_failure1(self):
|
||||
v1=Variable('v1', 0, 8)
|
||||
v2=Variable('v2', 0, 2)
|
||||
|
||||
+117
-96
@@ -3,7 +3,6 @@ import unittest, pickle, functools, math
|
||||
import z3
|
||||
|
||||
from tinygrad.dtype import dtypes, ConstType, DType, Invalid
|
||||
from tinygrad.helpers import Context
|
||||
from test.helpers import get_uops
|
||||
from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
|
||||
from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
|
||||
@@ -181,8 +180,8 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(Variable("a", 0, 8)*1, 0, 8, "a")
|
||||
|
||||
def test_mul_neg_1(self):
|
||||
self.helper_test_variable((Variable("a", 0, 2)*-1)//3, 0, 0, "0")
|
||||
self.helper_test_variable((Variable("a", 2, 7)*-1)//3, -2, 0, "((a//3)*-1)")
|
||||
self.helper_test_variable((Variable("a", 0, 2)*-1)//3, -1, 0, "((a*-1)//3)")
|
||||
self.helper_test_variable((Variable("a", 2, 7)*-1)//3, -3, -1, "((a*-1)//3)")
|
||||
|
||||
def test_mul_2(self):
|
||||
self.helper_test_variable(Variable("a", 0, 8)*2, 0, 16, "(a*2)")
|
||||
@@ -203,8 +202,8 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(Variable("a", 0, 7) // 20, 0, 0, "0")
|
||||
|
||||
def test_div_neg_min_max(self):
|
||||
self.helper_test_variable(Variable("a", 1, 7) // -2, -3, 0, "((a//2)*-1)")
|
||||
self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "((a//2)*-1)")
|
||||
self.helper_test_variable(Variable("a", 1, 7) // -2, -4, -1, "(a//-2)")
|
||||
self.helper_test_variable(Variable("a", 0, 6) // -2, -3, 0, "(a//-2)")
|
||||
|
||||
def test_div_mod_zero(self):
|
||||
with self.assertRaises(ZeroDivisionError):
|
||||
@@ -238,14 +237,14 @@ class TestSymbolic(unittest.TestCase):
|
||||
|
||||
def test_mod_min_max(self):
|
||||
self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
|
||||
self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), -9, 0, "(((x*-1)%y)*-1)")
|
||||
self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), 0, 9, "(x%(y*-1))")
|
||||
self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(((x*-1)%(y*-1))*-1)")
|
||||
self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 9, "(x%(y*-1))")
|
||||
self.helper_test_variable(Variable("x", -10, 0)%Variable("y", 1, 10), 0, 9, "(x%y)")
|
||||
self.helper_test_variable(Variable("x", 0, 10)%Variable("y", -10, -1), -9, 0, "(x%y)")
|
||||
self.helper_test_variable(Variable("x", -10, 0)%Variable("y", -10, -1), -9, 0, "(x%y)")
|
||||
self.helper_test_variable(Variable("x", -10, 10)%Variable("y", -10, -1), -9, 0, "(x%y)")
|
||||
|
||||
# test _min_max directly without the rewrite taking out the sign
|
||||
# test _min_max directly: floor mod with positive divisor is in [0, c-1]; with negative divisor in [c+1, 0]
|
||||
self.assertEqual((Variable("x", -10, 0)%Variable("y", -10, -1))._min_max, (-9, 0))
|
||||
self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
|
||||
self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (0, 9))
|
||||
|
||||
def test_range_div_its_symbolic_bound(self):
|
||||
a = Variable("a", 1, 10, dtypes.weakint)
|
||||
@@ -262,12 +261,12 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(Variable("a", 0, 6) // 2, 0, 3, "(a//2)")
|
||||
|
||||
self.helper_test_variable(Variable("x", 0, 10)//Variable("y", 1, 10), 0, 10, "(x//y)")
|
||||
self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(((x*-1)//y)*-1)")
|
||||
self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "((x//(y*-1))*-1)")
|
||||
self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "((x*-1)//(y*-1))")
|
||||
self.helper_test_variable(Variable("x", -10, 0)//Variable("y", 1, 10), -10, 0, "(x//y)")
|
||||
self.helper_test_variable(Variable("x", 0, 10)//Variable("y", -10, -1), -10, 0, "(x//y)")
|
||||
self.helper_test_variable(Variable("x", -10, 0)//Variable("y", -10, -1), 0, 10, "(x//y)")
|
||||
|
||||
self.helper_test_variable(Variable("x", -10, 10)//Variable("y", 1, 10), -10, 10, "(x//y)")
|
||||
self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "((x//(y*-1))*-1)")
|
||||
self.helper_test_variable(Variable("x", -10, 10)//Variable("y", -10, -1), -10, 10, "(x//y)")
|
||||
|
||||
def test_mod_factor(self):
|
||||
self.helper_test_variable(usum([Variable("a", 0, 7)*100, Variable("b", 0, 3)*50]) % 100, 0, 50, "((b%2)*50)")
|
||||
@@ -334,12 +333,12 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_mod_mod_wrong_sign(self):
|
||||
v1=Variable("v1", 0, 128)
|
||||
v3=Variable("v3", 0, 7)
|
||||
self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), -3, 4, "(v1%2*2+(v3+-1)%5+-2)")
|
||||
self.helper_test_variable((((((v1%2)*2)+((v3+-1)%5))+-2)%5), 0, 4, "((v3+v1%2*2+-3)%5)")
|
||||
|
||||
def test_mod_mod_wrong_sign2(self):
|
||||
v2=Variable("v2", 0, 8)
|
||||
v3=Variable("v3", 0, 4)
|
||||
self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), -2, 6, "(((v2+((v3+3)%7))+-2)%7)")
|
||||
self.helper_test_variable((((((v3+3)%7)+(v2+-2))%7)%7), 0, 6, "((v2+v3+1)%7)")
|
||||
|
||||
def test_mul_mul(self):
|
||||
self.helper_test_variable((Variable("a", 0, 5)*10)*9, 0, 5*10*9, "(a*90)")
|
||||
@@ -357,21 +356,28 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_div_const_div(self):
|
||||
a = Variable("a", 0, 124)
|
||||
self.helper_test_variable((a//2+1)//2, 0, 31, "((a+2)//4)")
|
||||
self.helper_test_variable(((-a)//2-1)//2, -31, 0, "(((a+2)//4)*-1)")
|
||||
self.helper_test_variable(((-a)//2+10)//2, -26, 5, "((((a//2)*-1)+10)//2)")
|
||||
self.helper_test_variable(((-a)//2-1)//2, -32, -1, "((a*-1+2)//4+-1)")
|
||||
self.helper_test_variable(((-a)//2+10)//2, -26, 5, "(a*-1//4+5)")
|
||||
|
||||
def test_div_const_div_wrong_sign(self):
|
||||
a = Variable("a", 0, 124)
|
||||
self.helper_test_variable(((a-10)//2+10)//2, 2, 33, "((((a+-10)//2)+10)//2)")
|
||||
self.helper_test_variable(((a-10)//2+10)//2, 2, 33, "((a+2)//4+2)")
|
||||
|
||||
def test_div_const_div_wrong_sign_divisor(self):
|
||||
a = Variable("a", 0, 124)
|
||||
self.helper_test_variable(((a+10)//-2+10)//-4, -1, 14, "(((((a//2)*-1)+5)//4)*-1)")
|
||||
self.helper_test_variable(((a+10)//-2+10)//-4, -2, 14, "(((a+10)//-2+10)//-4)")
|
||||
|
||||
def test_nested_div_negative_divisor(self):
|
||||
# (x//c1)//c2 -> x//(c1*c2) only when c2>0
|
||||
a = Variable("a", 0, 124)
|
||||
self.helper_test_variable((a//-2)//-3, 0, 20, "((a//-2)//-3)")
|
||||
self.helper_test_variable((a//2)//-3, -21, 0, "((a//2)//-3)")
|
||||
self.helper_test_variable((a//-2)//3, -21, 0, "(a//-6)")
|
||||
|
||||
def test_neg_mod(self):
|
||||
a = Variable("a", 0, 124)
|
||||
self.helper_test_variable((-a)%4, -3, 0, "((a%4)*-1)")
|
||||
self.helper_test_variable(a%-4, 0, 3, "(a%4)")
|
||||
self.helper_test_variable((-a)%4, 0, 3, "(a*-1%4)")
|
||||
self.helper_test_variable(a%-4, -3, 0, "(a%-4)")
|
||||
|
||||
def test_distribute_mul(self):
|
||||
self.helper_test_variable(usum([Variable("a", 0, 3), Variable("b", 0, 5)])*3, 0, 24, "((a*3)+(b*3))")
|
||||
@@ -387,11 +393,11 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(Variable("a", 0, 6)%100, 0, 6, "a")
|
||||
|
||||
def test_big_mod(self):
|
||||
self.helper_test_variable(Variable("a", -20, 20)%10, -9, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -20, 0)%10, -9, 0, "(((a*-1)%10)*-1)")
|
||||
self.helper_test_variable(Variable("a", -20, 1)%10, -9, 1, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -20, 20)%10, 0, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -20, 0)%10, 0, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -20, 1)%10, 0, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", 0, 20)%10, 0, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -1, 20)%10, -1, 9, "(a%10)")
|
||||
self.helper_test_variable(Variable("a", -1, 20)%10, 0, 9, "(a%10)")
|
||||
|
||||
def test_ge_remove(self):
|
||||
self.helper_test_variable(Variable("a", 0, 6) >= 25, 0, 0, "False")
|
||||
@@ -439,8 +445,8 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(c & c.logical_not(), False, False, "False")
|
||||
|
||||
def test_mod_factor_negative(self):
|
||||
self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)")
|
||||
self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, -27, 27, "(((a+(b*28))+-29)%28)")
|
||||
self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 10), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
|
||||
self.helper_test_variable(usum([uconst(-29), Variable("a", 0, 100), Variable("b", 0, 10)*28]) % 28, 0, 27, "((a+b*28+-29)%28)")
|
||||
|
||||
def test_sum_combine_num(self):
|
||||
self.helper_test_variable(usum([uconst(29), Variable("a", 0, 10), uconst(-23)]), 6, 16, "(a+6)")
|
||||
@@ -448,22 +454,12 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_sum_num_hoisted_and_factors_cancel_out(self):
|
||||
self.helper_test_variable(usum([Variable("a", 0, 1) * -4 + 1, Variable("a", 0, 1) * 4]), 1, 1, "1")
|
||||
|
||||
@unittest.expectedFailure # only correct for floordiv, not truncdiv
|
||||
def test_div_cancel(self):
|
||||
self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40])//40, -1, 9, "(b+-1)")
|
||||
|
||||
def test_div_cancel_correct(self):
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40])//40, -1, 9, "(((a+(b*20))+-20)//20)")
|
||||
|
||||
@unittest.expectedFailure # only correct for floordiv, not truncdiv
|
||||
def test_mod_cancel(self):
|
||||
self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) % 40, 0, 20, "(a*2)")
|
||||
|
||||
def test_mod_cancel_correct(self):
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
self.helper_test_variable(usum([uconst(-40), Variable("a", 0, 10)*2, Variable("b", 0, 10)*40]) % 40, -38, 38, "((((a+(b*20))+-20)%20)*2)")
|
||||
|
||||
def test_mul_div(self):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*4)//4, 0, 10, "a")
|
||||
|
||||
@@ -475,22 +471,22 @@ class TestSymbolic(unittest.TestCase):
|
||||
lidx1 = UOp.variable("lidx1", 0, 1)
|
||||
ridx1005 = UOp.variable("ridx1005", 0, 2)
|
||||
ridx1006 = UOp.variable("ridx1006", 0, 2)
|
||||
self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -2, 20,
|
||||
"(((((lidx1+(((gidx1*18)+(ridx1005*18))+(lidx0*162)))+(gidx0*2))+(ridx1006*2))+-40)//18)")
|
||||
self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
|
||||
"(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")
|
||||
|
||||
def test_add_div(self):
|
||||
# careful about the lower bounds and upper bounds
|
||||
self.helper_test_variable((Variable("a", 0, 5)-2)//4, 0, 0, "0")
|
||||
self.helper_test_variable((Variable("a", 0, 5)-1)//4, 0, 1, "((a+-1)//4)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)-2)//4, -1, 0, "((a+2)//4+-1)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)-1)//4, -1, 1, "((a+3)//4+-1)")
|
||||
self.helper_test_variable((Variable("a", 0, 5))//4, 0, 1, "(a//4)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+1)//4, 0, 1, "((a+1)//4)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+2)//4, 0, 1, "((a+2)//4)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+3)//4, 0, 2, "((a+3)//4)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+4)//4, 1, 2, "((a//4)+1)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+5)//4, 1, 2, "(((a+1)//4)+1)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+4)//4, 1, 2, "(a//4+1)")
|
||||
self.helper_test_variable((Variable("a", 0, 5)+5)//4, 1, 2, "((a+1)//4+1)")
|
||||
|
||||
def test_div_neg_rem(self):
|
||||
self.helper_test_variable((-Variable("a", 0, 255)+256)//2, 0, 128, "((((a+1)//2)*-1)+128)")
|
||||
self.helper_test_variable((-Variable("a", 0, 255)+256)//2, 0, 128, "(a*-1//2+128)")
|
||||
|
||||
def test_mul_div_factor_mul(self):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*8)//4, 0, 20, "(a*2)")
|
||||
@@ -502,7 +498,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*4)//8, 0, 5, "(a//2)")
|
||||
|
||||
def test_mul_div_factor_div_neg(self):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -4, 0, "(((a*-1)+1)//2)")
|
||||
self.helper_test_variable((Variable("a", 0, 10)*-4+4)//8, -5, 0, "((a*-1+1)//2)")
|
||||
|
||||
def test_div_symbolic_const_gcd(self):
|
||||
a = Variable("a", -10, 10)
|
||||
@@ -520,8 +516,8 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((d1*a*d2*b*d1)//(d1*d2), -1000, 1000, "(a*(b*d1))", test_z3=False)
|
||||
self.helper_test_variable((d1*a + b*d1)//(d1), -20, 20, "(a+b)", test_z3=False)
|
||||
self.helper_test_variable((d1*a + b*d1 + c*d1)//(d1), -30, 30, "(c+(a+b))", test_z3=False)
|
||||
self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "(((a+(b*3))//(d2*-1))*-1)", test_z3=False)
|
||||
self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "(((((a*d1)+((b*d1)*3))+1)//((d1*d2)*-1))*-1)", test_z3=False)
|
||||
self.helper_test_variable((3*a*d1 + 9*b*d1)//(3*d1*d2), -40, 40, "((a+b*3)//d2)", test_z3=False)
|
||||
self.helper_test_variable((3*a*d1 + 9*b*d1+3)//(3*d1*d2), -401, 399, "((a*d1+b*d1*3+1)//(d1*d2))", test_z3=False)
|
||||
|
||||
def test_symbolic_factor_remainder_div(self):
|
||||
a = Variable("a", 0, 10)
|
||||
@@ -532,7 +528,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((d*a*20+b*d*5+10)//(5*d), 0, 52, "((b+(a*4))+(2//d))")
|
||||
|
||||
def test_mod_gcd_factor_neg(self):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, -4, 4, "((((a*-1)+1)%2)*4)")
|
||||
self.helper_test_variable((Variable("a", 0, 10)*-4+4)%8, 0, 4, "((a*-1+1)%2*4)")
|
||||
|
||||
def test_mod_gcd_fold_neg(self):
|
||||
self.helper_test_variable((Variable("a", 0, 10)*-8+20)%4, 0, 0, "0")
|
||||
@@ -540,22 +536,32 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_sum_div_partial_remove(self):
|
||||
self.helper_test_variable(usum([Variable("idx0", 0, 127)*4, Variable("idx2", 0, 3)])//4, 0, 127, "idx0")
|
||||
|
||||
def test_cdiv_const_evaluation(self):
|
||||
self.helper_test_variable((Variable("a", 0, 2)-12)//8, -1, -1, "-1")
|
||||
self.helper_test_variable((-Variable("a", 0, 2))//7, 0, 0, "0")
|
||||
def test_floordiv_const_evaluation(self):
|
||||
self.helper_test_variable((Variable("a", 0, 2)-12)//8, -2, -2, "-2")
|
||||
self.helper_test_variable((-Variable("a", 0, 2))//7, -1, 0, "(a*-1//7)")
|
||||
|
||||
def test_cmod_const_evaluation(self):
|
||||
self.helper_test_variable((Variable("a", 1, 1)*-3)%8, -3, -3, "-3")
|
||||
self.helper_test_variable((-Variable("a", 10, 10))%7, -3, -3, "-3")
|
||||
def test_floormod_const_evaluation(self):
|
||||
self.helper_test_variable((Variable("a", 1, 1)*-3)%8, 5, 5, "5")
|
||||
self.helper_test_variable((-Variable("a", 10, 10))%7, 4, 4, "4")
|
||||
|
||||
def test_div_numerator_negative(self):
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -8, 0, "(((idx*10)//11)*-1)")
|
||||
self.helper_test_variable((Variable("idx", 0, 9)*-10)//11, -9, 0, "(idx*-1)")
|
||||
|
||||
def test_nest_div_negative_factor(self):
|
||||
ridx0=Variable("ridx0", 0, 9)
|
||||
ridx1=Variable("ridx1", 0, 6)
|
||||
self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "(((ridx0//5)*-1)+1)")
|
||||
self.helper_test_variable(((((ridx0*-7)+ridx1)+63)//35), 0, 1, "((ridx0*-1+4)//5+1)")
|
||||
|
||||
def test_floordiv_factor_nest_negative_numerator(self):
|
||||
# x//c = (x//f)//(c//f) for f|c, any sign of x
|
||||
a = Variable("a", -10, 10)
|
||||
b = Variable("b", 0, 3)
|
||||
self.helper_test_variable((a*4 + b)//12, -4, 3, "(a//3)")
|
||||
|
||||
def test_floordiv_gcd_with_remainder_negative_numerator(self):
|
||||
# factor gcd from numerator, even when x crosses zero, as long as the shifted numerator stays nonneg
|
||||
a = Variable("a", -1, 5)
|
||||
self.helper_test_variable((a*2 + 7)//8, 0, 2, "((a+3)//4)")
|
||||
|
||||
def test_div_into_mod(self):
|
||||
self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")
|
||||
@@ -568,11 +574,11 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable(x%12//4*4 + x%4 + x//12*12, 0, 23, "x")
|
||||
|
||||
def test_div_neg_cancel(self):
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 1, 26, "((idx//4)+1)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx+3)//4)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, 0, 25, "((idx+2)//4)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "((idx//2)*-1)")
|
||||
self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "((idx//2)*-1)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 0, 25, "((idx*-1+199)//-4+50)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx*-1+200)//-4+50)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100)+201)//-4 + 50, -1, 24, "((idx*-1+201)//-4+50)")
|
||||
self.helper_test_variable((-Variable("idx", 0, 100))//2, -50, 0, "(idx*-1//2)")
|
||||
self.helper_test_variable(Variable("idx", 0, 100)//-2, -50, 0, "(idx//-2)")
|
||||
|
||||
def test_sum_div_big_const(self):
|
||||
gidx0 = Variable("gidx0", 0, 24)
|
||||
@@ -647,22 +653,22 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_div_neg_all_range(self):
|
||||
gidx = Variable("gidx", 0, 124)
|
||||
lidx = Variable("lidx", 0, 7)
|
||||
self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 1, 250, "(((gidx*2)+(lidx//4))+1)")
|
||||
self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 250, "((gidx*2)+((lidx+3)//4))")
|
||||
self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, 0, 250, "((gidx*2)+((lidx+2)//4))")
|
||||
self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, 0, 250, "((gidx*2)+((lidx+1)//4))")
|
||||
self.helper_test_variable((-gidx*8-lidx+999)//-4 + 250, 0, 250, "((gidx*-8+lidx*-1+999)//-4+250)")
|
||||
self.helper_test_variable((-gidx*8-lidx+1000)//-4 + 250, 0, 249, "((gidx*-8+lidx*-1+1000)//-4+250)")
|
||||
self.helper_test_variable((-gidx*8-lidx+1001)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1001)//-4+250)")
|
||||
self.helper_test_variable((-gidx*8-lidx+1002)//-4 + 250, -1, 249, "((gidx*-8+lidx*-1+1002)//-4+250)")
|
||||
|
||||
def test_div_neg_then_neg(self):
|
||||
# taken from arange opts
|
||||
lidx0 = Variable("lidx0", 0, 7)
|
||||
lidx1 = Variable("lidx1", 0, 7)
|
||||
alu2 = -lidx0-lidx1
|
||||
self.helper_test_variable((((alu2+14)//(-32))+4), 4, 4, "4")
|
||||
self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -4, "-4")
|
||||
self.helper_test_variable((((alu2+134)//(-32))+4), 0, 1, "(((lidx0+lidx1)+25)//32)")
|
||||
self.helper_test_variable((((alu2+142)//(-32))+4), 0, 0, "0")
|
||||
self.helper_test_variable((((alu2+150)//(-32))+4), 0, 0, "0")
|
||||
self.helper_test_variable((((alu2+158)//(-32))+4), 0, 0, "0")
|
||||
self.helper_test_variable((((alu2+14)//(-32))+4), 3, 4, "((lidx0*-1+lidx1*-1+14)//-32+4)")
|
||||
self.helper_test_variable(-(((alu2+14)//(-32))+4), -4, -3, "((lidx0*-1+lidx1*-1+14)//-32*-1+-4)")
|
||||
self.helper_test_variable((((alu2+134)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+134)//-32+4)")
|
||||
self.helper_test_variable((((alu2+142)//(-32))+4), -1, 0, "((lidx0*-1+lidx1*-1+142)//-32+4)")
|
||||
self.helper_test_variable((((alu2+150)//(-32))+4), -1, -1, "-1")
|
||||
self.helper_test_variable((((alu2+158)//(-32))+4), -1, -1, "-1")
|
||||
|
||||
def test_div_mod_recombine(self):
|
||||
gidx = Variable("gidx", 0, 124)
|
||||
@@ -696,7 +702,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
# negative variable range
|
||||
xn = Variable("x", -1000, 1000)
|
||||
self.helper_test_variable(xn//3%224*3 + xn%3 + xn//672*672, -1000, 1000, "x")
|
||||
self.helper_test_variable(xn//3%7*3 + xn//21*21, -999, 999, "(x//3*3)")
|
||||
self.helper_test_variable(xn//3%7*3 + xn//21*21, -1002, 999, "(x//3*3)")
|
||||
# should NOT simplify: a*c1 != b (3*224 != 600)
|
||||
self.helper_test_variable(gidx//3%224*3 + gidx//600*600, 0, 150669, "(gidx//600*600+gidx//3%224*3)")
|
||||
# should NOT simplify: c1*c2 != c3 (224*3 != 700)
|
||||
@@ -709,7 +715,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")
|
||||
|
||||
def test_div_partial_quotient(self):
|
||||
# IDIV should extract partial quotients when const_factor > divisor, matching what MOD already does
|
||||
# FLOORDIV should extract partial quotients when const_factor > divisor, matching what FLOORMOD already does
|
||||
# (f*x+c)//d -> (f%d*x+c)//d + (f//d)*x when f >= d
|
||||
b = Variable("b", 0, 100)
|
||||
self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
|
||||
@@ -730,8 +736,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
def test_div_by_factor_tie_break(self):
|
||||
a = Variable("a", 0, 1)
|
||||
b = Variable("b", 0, 1)
|
||||
with Context(CORRECT_DIVMOD_FOLDING=1):
|
||||
self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
|
||||
self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
|
||||
|
||||
def test_div_mod_recombine_large_coeff(self):
|
||||
# recombine must work even when coeff > divisor: both mod and div reduce the coeff the same way
|
||||
@@ -741,7 +746,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
|
||||
|
||||
def test_mod_nest_by_factor(self):
|
||||
# (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for MOD
|
||||
# (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for FLOORMOD
|
||||
gidx0 = Variable("gidx0", 0, 15)
|
||||
lidx0 = Variable("lidx0", 0, 3)
|
||||
# f=4, k=2, c=8: (gidx0*4+lidx0)%8 = (gidx0%2)*4 + lidx0
|
||||
@@ -755,7 +760,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((a*3+b)%9, 0, 8, "(b+a%3*3)")
|
||||
|
||||
def test_mod_nest_by_factor_with_const(self):
|
||||
# nest_by_factor MOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
|
||||
# nest_by_factor FLOORMOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
|
||||
a = Variable("a", 0, 7)
|
||||
b = Variable("b", 0, 1)
|
||||
# f=4, k=2, const=2: (a*4+b+2)%8 = (a%2)*4 + b + 2
|
||||
@@ -767,7 +772,7 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
|
||||
|
||||
def test_div_nest_by_factor_with_const(self):
|
||||
# nest_by_factor IDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
|
||||
# nest_by_factor FLOORDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
|
||||
a = Variable("a", 0, 2)
|
||||
b = Variable("b", 0, 31)
|
||||
c = Variable("c", 0, 1)
|
||||
@@ -827,12 +832,26 @@ class TestSymbolic(unittest.TestCase):
|
||||
# TODO: simplify the true branch
|
||||
self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
|
||||
|
||||
def test_idiv_lt(self):
|
||||
def test_floordiv_lt(self):
|
||||
# x//d<c <=> x<c*d for d>0
|
||||
idx = Variable("idx", 0, 24)
|
||||
self.helper_test_variable((idx//4<3), 0, 1, "(idx<12)")
|
||||
self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<5)")
|
||||
self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<7)")
|
||||
self.helper_test_variable((idx//-4<-3), 0, 1, "(((idx//4)*-1)<-3)")
|
||||
self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<8)")
|
||||
self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<10)")
|
||||
self.helper_test_variable((idx//-4<-3), 0, 1, "((idx//-4)<-3)")
|
||||
|
||||
def test_nested_div_mod_negative_inner_divisor(self):
|
||||
# (x % (k*c)) // c -> (x // c) % k requires k>0; (x % (k*c)) % c -> x % c is unconditional for c>0
|
||||
a = Variable("a", 0, 100)
|
||||
self.helper_test_variable((a % -8) // 2, -4, 0, "(a%-8//2)")
|
||||
self.helper_test_variable((a % -8) % 2, 0, 1, "(a%2)")
|
||||
|
||||
def test_floordiv_lt_negative_c(self):
|
||||
# x//d<c with negative c also reduces to x<c*d for d>0
|
||||
idx = Variable("idx", -20, 20)
|
||||
self.helper_test_variable((idx//4 < 0), 0, 1, "(idx<0)")
|
||||
self.helper_test_variable((idx//4 < -1), 0, 1, "(idx<-4)")
|
||||
self.helper_test_variable((idx//4 < -2), 0, 1, "(idx<-8)")
|
||||
|
||||
def test_simplex_lt(self):
|
||||
a = Variable("a", 0, 3)
|
||||
@@ -981,10 +1000,10 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")
|
||||
|
||||
def test_nested_mod_negative_range(self):
|
||||
# (x%(k*c))%c = x%c holds for cmod regardless of signs since sign(x%(k*c)) = sign(x)
|
||||
# (x%(k*c))%c = x%c for positive c
|
||||
x = Variable("x", 0, 1575)
|
||||
self.helper_test_variable(((x + (-1064)) % 512) % 4, -3, 3, "((x+-1064)%4)")
|
||||
self.helper_test_variable(((x + (-1064)) % 512) % 128, -127, 127, "((x+-1064)%128)")
|
||||
self.helper_test_variable(((x + (-1064)) % 512) % 4, 0, 3, "((x+-1064)%4)")
|
||||
self.helper_test_variable(((x + (-1064)) % 512) % 128, 0, 127, "((x+-1064)%128)")
|
||||
|
||||
class TestSymbolicNumeric(unittest.TestCase):
|
||||
def helper_test_numeric(self, f):
|
||||
@@ -1062,12 +1081,13 @@ class TestSymInfer(unittest.TestCase):
|
||||
assert sym_infer(a+b+c, var_vals) == 9
|
||||
assert sym_infer(a*b, var_vals) == 6
|
||||
assert sym_infer(a*b+c, var_vals) == 10
|
||||
def test_sym_infer_cdiv_cmod(self):
|
||||
def test_sym_infer_floordiv_floormod(self):
|
||||
a = Variable("a", -1000, 1)
|
||||
b = Variable("b", -1000, 1)
|
||||
var_vals = {a.expr: 1, b.expr: -1000}
|
||||
assert sym_infer(a%b, var_vals) == 1
|
||||
assert sym_infer(a//b, var_vals) == 0
|
||||
# floor: 1 % -1000 = -999, 1 // -1000 = -1
|
||||
assert sym_infer(a%b, var_vals) == -999
|
||||
assert sym_infer(a//b, var_vals) == -1
|
||||
def test_sym_infer_with_bitcast(self):
|
||||
a = Variable("a", 1, 10, dtypes.int)
|
||||
expr = ((a.bitcast(dtypes.uint) << UOp.const(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
|
||||
@@ -1286,7 +1306,8 @@ class TestGatedUopGivenValid(unittest.TestCase):
|
||||
|
||||
idx:UOp = (r0 < 3).where((r0 + uconst(-1)) // uconst(3), UOp.invalid())
|
||||
idx = graph_rewrite(idx, pm_simplify_valid)
|
||||
self.assertEqual(idx, (r0 < 3).where(uconst(0), UOp.invalid()))
|
||||
# (r0-1)//3 = (r0+2)//3 - 1 (constant offset split)
|
||||
self.assertEqual(idx, (r0 < 3).where((r0 + uconst(2)) // uconst(3) + uconst(-1), UOp.invalid()))
|
||||
|
||||
def test_invalid_gate_simplifies_vectorize(self):
|
||||
r0 = Variable("r0", 0, 2)
|
||||
@@ -1295,8 +1316,8 @@ class TestGatedUopGivenValid(unittest.TestCase):
|
||||
idx1 = r0 % uconst(3)
|
||||
idx:UOp = (r0 < 3).where(UOp(Ops.STACK, dtypes.weakint.vec(2), (idx0, idx1)), UOp.invalid())
|
||||
idx = graph_rewrite(idx, pm_simplify_valid)
|
||||
# NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
|
||||
expected_vec = UOp(Ops.STACK, dtypes.weakint.vec(2), (uconst(0), r0))
|
||||
# independent simplification: (r0-1)//3 -> (r0+2)//3 - 1, and r0%3 -> r0 when r0 in [0,2]
|
||||
expected_vec = UOp(Ops.STACK, dtypes.weakint.vec(2), ((r0 + uconst(2)) // uconst(3) + uconst(-1), r0))
|
||||
self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
|
||||
|
||||
class TestRangeSplitting(unittest.TestCase):
|
||||
@@ -1335,8 +1356,8 @@ class TestBounds(unittest.TestCase):
|
||||
alu0 = gidx0 * -1
|
||||
assert alu0.vmin == -2559 and alu0.vmax == 0
|
||||
assert (alu0+2559).vmin == 0 and (alu0+2559).vmax == 2559
|
||||
assert ((alu0+2559)//-4).vmin == -639 and ((alu0+2559)//-4).vmax == 0
|
||||
assert (((alu0+2559)//-4)*(-1)).vmin == 0 and (((alu0+2559)//-4)*(-1)).vmax == 639
|
||||
assert ((alu0+2559)//-4).vmin == -640 and ((alu0+2559)//-4).vmax == 0
|
||||
assert (((alu0+2559)//-4)*(-1)).vmin == 0 and (((alu0+2559)//-4)*(-1)).vmax == 640
|
||||
|
||||
class TestFuzzFailure(unittest.TestCase):
|
||||
def test_fuzz_failure1(self):
|
||||
|
||||
@@ -173,17 +173,15 @@ class TestVminVmaxDivMod(unittest.TestCase):
|
||||
self.assertEqual(uop.vmax, 10)
|
||||
|
||||
def test_vmin_vmax_division_negative(self):
|
||||
# vmin and vmax for division of a variable by a negative constant
|
||||
# always positive
|
||||
# floor division of a variable by a negative constant
|
||||
x = UOp.variable('x', 10, 20)
|
||||
uop = x // -2
|
||||
self.assertEqual(uop.vmin, -10)
|
||||
self.assertEqual(uop.vmax, -5)
|
||||
uop = x // -3
|
||||
self.assertEqual(uop.vmin, -6)
|
||||
self.assertEqual(uop.vmax, -3)
|
||||
self.assertEqual(uop.vmin, -7)
|
||||
self.assertEqual(uop.vmax, -4)
|
||||
|
||||
# always negative
|
||||
x = UOp.variable('x', -20, -10)
|
||||
uop = x // -2
|
||||
self.assertEqual(uop.vmin, 5)
|
||||
@@ -193,7 +191,6 @@ class TestVminVmaxDivMod(unittest.TestCase):
|
||||
self.assertEqual(uop.vmax, 6)
|
||||
|
||||
def test_vmin_vmax_floordiv_floormod(self):
|
||||
# FLOORDIV/FLOORMOD ranges differ from IDIV/MOD when the dividend can be negative
|
||||
x = UOp.variable('x', -7, 7)
|
||||
floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
|
||||
self.assertEqual(floordiv.vmin, -3)
|
||||
@@ -212,32 +209,42 @@ class TestVminVmaxDivMod(unittest.TestCase):
|
||||
self.assertEqual(uop.vmin, -5)
|
||||
self.assertEqual(uop.vmax, 5)
|
||||
uop = x // -3
|
||||
self.assertEqual(uop.vmin, -3)
|
||||
self.assertEqual(uop.vmin, -4)
|
||||
self.assertEqual(uop.vmax, 3)
|
||||
|
||||
def test_vmin_vmax_floordiv_floormod_empty_range(self):
|
||||
# empty numerator range (vmin > vmax, e.g. RANGE with end=0) short-circuits to (0, 0)
|
||||
rng = UOp.range(0, 0)
|
||||
self.assertEqual(rng.vmin, 0)
|
||||
self.assertEqual(rng.vmax, -1)
|
||||
self.assertEqual((rng // 4).vmin, 0)
|
||||
self.assertEqual((rng // 4).vmax, 0)
|
||||
self.assertEqual((rng % 4).vmin, 0)
|
||||
self.assertEqual((rng % 4).vmax, 0)
|
||||
|
||||
def test_vmin_vmax_div_symbolic(self):
|
||||
x = UOp.variable('x', 1, 10)
|
||||
y = UOp.variable('y', 3, 5)
|
||||
self.assertEqual((x//y).vmin, 0)
|
||||
self.assertEqual((x//y).vmax, 3)
|
||||
self.assertEqual(((-x)//y).vmin, -3)
|
||||
self.assertEqual(((-x)//y).vmax, 0)
|
||||
self.assertEqual((x//(-y)).vmin, -3)
|
||||
self.assertEqual((x//(-y)).vmax, 0)
|
||||
self.assertEqual(((-x)//y).vmin, -4)
|
||||
self.assertEqual(((-x)//y).vmax, -1)
|
||||
self.assertEqual((x//(-y)).vmin, -4)
|
||||
self.assertEqual((x//(-y)).vmax, -1)
|
||||
self.assertEqual(((-x)//(-y)).vmin, 0)
|
||||
self.assertEqual(((-x)//(-y)).vmax, 3)
|
||||
|
||||
self.assertEqual((100//y).vmin, 20)
|
||||
self.assertEqual((100//y).vmax, 33)
|
||||
self.assertEqual(((-100)//y).vmin, -33)
|
||||
self.assertEqual(((-100)//y).vmin, -34)
|
||||
self.assertEqual(((-100)//y).vmax, -20)
|
||||
self.assertEqual((100//(-y)).vmin, -33)
|
||||
self.assertEqual((100//(-y)).vmin, -34)
|
||||
self.assertEqual((100//(-y)).vmax, -20)
|
||||
self.assertEqual(((-100)//(-y)).vmin, 20)
|
||||
self.assertEqual(((-100)//(-y)).vmax, 33)
|
||||
|
||||
def test_vmin_vmax_mod_positive(self):
|
||||
# vmin and vmax for modulo of a variable by a positive constant
|
||||
# floor mod with positive divisor: result in [0, c-1] regardless of dividend sign
|
||||
positive = UOp.variable('positive', 10, 20)
|
||||
uop = positive % 3
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
@@ -245,20 +252,20 @@ class TestVminVmaxDivMod(unittest.TestCase):
|
||||
|
||||
negative = UOp.variable('negative', -20, -10)
|
||||
uop = negative % 3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
|
||||
mixed = UOp.variable('mixed', -20, 20)
|
||||
uop = mixed % 3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
|
||||
def test_vmin_vmax_mod_negative(self):
|
||||
# vmin and vmax for modulo of a variable by a negative constant
|
||||
# floor mod with negative divisor: result in [c+1, 0] regardless of dividend sign
|
||||
positive = UOp.variable('positive', 10, 20)
|
||||
uop = positive % -3
|
||||
self.assertEqual(uop.vmin, 0)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
|
||||
negative = UOp.variable('negative', -20, -10)
|
||||
uop = negative % -3
|
||||
@@ -268,7 +275,7 @@ class TestVminVmaxDivMod(unittest.TestCase):
|
||||
mixed = UOp.variable('mixed', -20, 20)
|
||||
uop = mixed % -3
|
||||
self.assertEqual(uop.vmin, -2)
|
||||
self.assertEqual(uop.vmax, 2)
|
||||
self.assertEqual(uop.vmax, 0)
|
||||
|
||||
class TestVminVmaxVConst(unittest.TestCase):
|
||||
def test_vmin_vmax_vconst_single_element(self):
|
||||
|
||||
@@ -177,6 +177,30 @@ class TestFastIdiv(unittest.TestCase):
|
||||
self.assertIn(Ops.SHR, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.IDIV, ops, f"For dtype={dt} divison by power of two did not simplify to shift")
|
||||
|
||||
def test_floormod_power_of_two(self):
|
||||
# FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
|
||||
for dt in (dtypes.int32, dtypes.uint32):
|
||||
g = UOp(Ops.PARAM, dt.ptr(), (), 0)
|
||||
c = UOp.const(dt, 8)
|
||||
a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
|
||||
self.assertNotIn(Ops.MOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
|
||||
self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
|
||||
|
||||
def test_floordiv_power_of_two_uint(self):
|
||||
# uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
|
||||
for dt in (dtypes.uint32, dtypes.uint64):
|
||||
g = UOp(Ops.PARAM, dt.ptr(), (), 0)
|
||||
c = UOp.const(dt, 2)
|
||||
a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
|
||||
uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
|
||||
ops = [x.op for x in uops]
|
||||
self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.IDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
|
||||
self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
|
||||
def test_fast_idiv_and_mod(self):
|
||||
g = UOp(Ops.PARAM, dtypes.uint32.ptr(), (), 0)
|
||||
|
||||
@@ -411,11 +411,11 @@ class TestVizIntegration(unittest.TestCase):
|
||||
def test_jit(self):
|
||||
with save_viz():
|
||||
@TinyJit
|
||||
def f(a, b, c): return (a+b).contiguous().mul(3), c.assign(a.to(c.device))
|
||||
def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
|
||||
a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
|
||||
for _ in range(3): Tensor.realize(*f(a, b, c))
|
||||
out = load_profile(cpu_events)
|
||||
self.assertEqual(["NULL", "NULL Graph", "NULL:SDMA:0"], [k for k in out["layout"] if k.startswith("NULL")])
|
||||
self.assertEqual(["NULL", "NULL Graph", "NULL:SDMA:0", "NULL:1", "NULL:1:SDMA:0"], [k for k in out["layout"] if k.startswith("NULL")])
|
||||
self.assertEqual(len(out["layout"]["NULL"]["events"]), 2*3)
|
||||
self.assertEqual(len(out["layout"]["NULL:SDMA:0"]["events"]), 3)
|
||||
self.assertEqual(len(out["layout"]["NULL Graph"]["events"]), 2)
|
||||
|
||||
@@ -17,7 +17,8 @@ pm_flatten_range = PatternMatcher([
|
||||
(UPat((Ops.REDUCE, Ops.END), name="r"), flatten_range),
|
||||
])
|
||||
|
||||
def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice)
|
||||
# index/range arithmetic uses FLOORDIV/FLOORMOD prior to late rewrite
|
||||
def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.FLOORDIV, Ops.FLOORMOD} for u in x.backward_slice)
|
||||
def simplify_merge_adjacent(u:UOp) -> UOp|None:
|
||||
reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
|
||||
# on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
|
||||
|
||||
+1
-1
@@ -241,7 +241,7 @@ SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), Contex
|
||||
RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
|
||||
CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)
|
||||
VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
|
||||
CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
|
||||
FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)
|
||||
ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
|
||||
MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0)
|
||||
EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
|
||||
|
||||
@@ -181,7 +181,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
|
||||
return self._binop(Ops.IDIV, x, reverse)
|
||||
|
||||
def mod(self, x: Self | ConstType, reverse: bool = False) -> Self:
|
||||
return self._binop(Ops.MOD, x, reverse)
|
||||
return self._binop(Ops.FLOORMOD, x, reverse)
|
||||
|
||||
def div(self, x: Self | ConstType, reverse: bool = False) -> Self:
|
||||
lhs, rhs = self._broadcasted(x, reverse)
|
||||
@@ -206,7 +206,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
|
||||
return self.div(x)
|
||||
|
||||
def __floordiv__(self, x: Self | ConstType) -> Self:
|
||||
return self.idiv(x) # TODO: idiv is trunc div, not floordiv
|
||||
return self._binop(Ops.FLOORDIV, x, False)
|
||||
|
||||
def __mod__(self, x: Self | ConstType) -> Self:
|
||||
return self.mod(x)
|
||||
@@ -233,7 +233,7 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
|
||||
return self.div(x, True)
|
||||
|
||||
def __rfloordiv__(self, x: Self | ConstType) -> Self:
|
||||
return self.idiv(x, True)
|
||||
return self._binop(Ops.FLOORDIV, x, True)
|
||||
|
||||
def __rand__(self, x: Self | ConstType) -> Self:
|
||||
return self.bitwise_and(x, True)
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
import pathlib, hashlib, re, itertools
|
||||
from tinygrad.runtime.autogen import load, root
|
||||
|
||||
__all__ = ["am", "pm4_soc15", "pm4_nv", "sdma_4_0_0", "sdma_5_0_0", "sdma_6_0_0", "smu_13_0_0", "smu_13_0_6", "smu_13_0_12", "smu_14_0_2",
|
||||
"fw", "navi_offsets", "vega_offsets", "regs", "soc_9", "soc_11", "soc_12"]
|
||||
|
||||
am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/33970e1351f5e511029602454979f3de7e22260f.tar.gz"
|
||||
rocm_src="https://github.com/ROCm/rocm-systems/archive/cccc350dc620e61ae2554978b62ab3532dc10bd9.tar.gz"
|
||||
AMD, AMDINC = "{}/drivers/gpu/drm/amd", "{}/drivers/gpu/drm/amd/include"
|
||||
inc, kern_rules = ["-include", "stdint.h"], [(r'le32_to_cpu', ''),]
|
||||
fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/1e2c15348485939baf1b6d1f5a7a3b799d80703d/1e2c15348485939baf1b6d1f5a7a3b799d80703d.tar.gz"
|
||||
@@ -26,6 +30,8 @@ reg_patterns = {
|
||||
"mp": ["MP([01]|ASP)_SMN_C2PMSG"], "hdp": ["HDP_MEM_POWER_CTRL"], "oss": ["IH_"], "sdma": ["SDMA_GFX", "SDMA_CNTL"]
|
||||
}
|
||||
|
||||
soc_patterns = ["SQ_TT", "VGT_EVENT_TYPE", "CS", "MTYPE", "SH"]
|
||||
|
||||
def __getattr__(nm):
|
||||
match nm:
|
||||
case "am": return load("am/am", [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
|
||||
@@ -40,13 +46,13 @@ def __getattr__(nm):
|
||||
args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
|
||||
case "sdma_6_0_0": return load("am/sdma_6_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/sdma_v6_0_0_pkt_open.h"],
|
||||
args=["-I/opt/rocm/include", "-x", "c++"], srcs=am_src)
|
||||
case "smu_v13_0_0": return load("am/smu_v13_0_0", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
|
||||
case "smu_13_0_0": return load("am/smu_13_0_0", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_0_ppsmc","smu13_driver_if_v13_0_0"]]
|
||||
+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
|
||||
case "smu_v13_0_6": return load("am/smu_v13_0_6", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
|
||||
case "smu_13_0_6": return load("am/smu_13_0_6", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_6_ppsmc","smu_v13_0_6_pmfw", \
|
||||
"smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
|
||||
case "smu_v13_0_12": return load("am/smu_v13_0_12", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
|
||||
case "smu_13_0_12": return load("am/smu_13_0_12", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v13_0_12_ppsmc","smu_v13_0_12_pmfw",
|
||||
"smu13_driver_if_v13_0_6"]] +[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
|
||||
case "smu_v14_0_2": return load("am/smu_v14_0_2", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
|
||||
case "smu_14_0_2": return load("am/smu_14_0_2", [f"{AMD}/pm/swsmu/inc/pmfw_if/{s}.h" for s in ["smu_v14_0_0_pmfw", "smu_v14_0_2_ppsmc",
|
||||
"smu14_driver_if_v14_0"]]+[root/"extra/amdpci/headers/amdgpu_smu.h"], args=inc, srcs=am_src)
|
||||
# firmware hashes
|
||||
case "fw":
|
||||
@@ -79,4 +85,7 @@ def __getattr__(nm):
|
||||
return "\n".join(out)
|
||||
return load("am/regs", [AMDINC + "/asic_reg/" + {"osssys":"oss"}.get(pre, pre) + f"/{pre}_{'_'.join(map(str, ver))}"
|
||||
for pre in reg_files for ver in sorted(reg_files[pre])], srcs=am_src, gen=genreg)
|
||||
case "soc_9" | "soc_11" | "soc_12":
|
||||
return load(f"am/{nm}", ["{}/projects/aqlprofile/linux/" + {9: "vega10", 11: "soc21", 12: "soc24"}[int(nm.split('_')[1])] + "_enum.h"],
|
||||
srcs=rocm_src, patterns=soc_patterns, macros=False)
|
||||
case _: raise AttributeError(f"no such autogen: {nm}")
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
# mypy: disable-error-code="empty-body"
|
||||
from __future__ import annotations
|
||||
import ctypes
|
||||
from typing import Literal, TypeAlias
|
||||
from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
|
||||
from tinygrad.runtime.support import c
|
||||
enum_MTYPE: dict[int, str] = {(MTYPE_C_RW_US:=0): 'MTYPE_C_RW_US', (MTYPE_RESERVED_1:=1): 'MTYPE_RESERVED_1', (MTYPE_C_RO_S:=2): 'MTYPE_C_RO_S', (MTYPE_UC:=3): 'MTYPE_UC', (MTYPE_C_RW_S:=4): 'MTYPE_C_RW_S', (MTYPE_RESERVED_5:=5): 'MTYPE_RESERVED_5', (MTYPE_C_RO_US:=6): 'MTYPE_C_RO_US', (MTYPE_RESERVED_7:=7): 'MTYPE_RESERVED_7'}
|
||||
MTYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
|
||||
SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
|
||||
SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_MODE: dict[int, str] = {(SQ_TT_MODE_OFF:=0): 'SQ_TT_MODE_OFF', (SQ_TT_MODE_ON:=1): 'SQ_TT_MODE_ON', (SQ_TT_MODE_GLOBAL:=2): 'SQ_TT_MODE_GLOBAL', (SQ_TT_MODE_DETAIL:=3): 'SQ_TT_MODE_DETAIL'}
|
||||
SQ_TT_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_RT_FREQ: dict[int, str] = {(SQ_TT_RT_FREQ_NEVER:=0): 'SQ_TT_RT_FREQ_NEVER', (SQ_TT_RT_FREQ_1024_CLK:=1): 'SQ_TT_RT_FREQ_1024_CLK', (SQ_TT_RT_FREQ_4096_CLK:=2): 'SQ_TT_RT_FREQ_4096_CLK'}
|
||||
SQ_TT_RT_FREQ: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_INST_EXCLUDE: dict[int, str] = {(SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_BIT:=1): 'SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_BIT', (SQ_TT_INST_EXCLUDE_EXPGNT234_BIT:=2): 'SQ_TT_INST_EXCLUDE_EXPGNT234_BIT'}
|
||||
SQ_TT_TOKEN_MASK_INST_EXCLUDE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_INST_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_SHIFT:=0): 'SQ_TT_INST_EXCLUDE_VMEM_OTHER_SIMD_SHIFT', (SQ_TT_INST_EXCLUDE_EXPGNT234_SHIFT:=1): 'SQ_TT_INST_EXCLUDE_EXPGNT234_SHIFT'}
|
||||
SQ_TT_TOKEN_MASK_INST_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_REG_EXCLUDE: dict[int, str] = {(SQ_TT_REG_EXCLUDE_USER_DATA_BIT:=1): 'SQ_TT_REG_EXCLUDE_USER_DATA_BIT', (SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_BIT:=2): 'SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_BIT', (SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_BIT:=4): 'SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_BIT'}
|
||||
SQ_TT_TOKEN_MASK_REG_EXCLUDE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_REG_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_REG_EXCLUDE_USER_DATA_SHIFT:=0): 'SQ_TT_REG_EXCLUDE_USER_DATA_SHIFT', (SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_SHIFT:=1): 'SQ_TT_REG_EXCLUDE_CP_ME_MC_RADDR_SHIFT', (SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_SHIFT:=2): 'SQ_TT_REG_EXCLUDE_GRBM_COMPUTE_EXCLUDE_SHIFT'}
|
||||
SQ_TT_TOKEN_MASK_REG_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_REG_INCLUDE: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_BIT:=1): 'SQ_TT_TOKEN_MASK_SQDEC_BIT', (SQ_TT_TOKEN_MASK_SHDEC_BIT:=2): 'SQ_TT_TOKEN_MASK_SHDEC_BIT', (SQ_TT_TOKEN_MASK_GFXUDEC_BIT:=4): 'SQ_TT_TOKEN_MASK_GFXUDEC_BIT', (SQ_TT_TOKEN_MASK_COMP_BIT:=8): 'SQ_TT_TOKEN_MASK_COMP_BIT', (SQ_TT_TOKEN_MASK_CONTEXT_BIT:=16): 'SQ_TT_TOKEN_MASK_CONTEXT_BIT', (SQ_TT_TOKEN_MASK_CONFIG_BIT:=32): 'SQ_TT_TOKEN_MASK_CONFIG_BIT', (SQ_TT_TOKEN_MASK_ALL_BIT:=64): 'SQ_TT_TOKEN_MASK_ALL_BIT', (SQ_TT_TOKEN_MASK_RSVD_BIT:=128): 'SQ_TT_TOKEN_MASK_RSVD_BIT'}
|
||||
SQ_TT_TOKEN_MASK_REG_INCLUDE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_REG_INCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_SHIFT:=0): 'SQ_TT_TOKEN_MASK_SQDEC_SHIFT', (SQ_TT_TOKEN_MASK_SHDEC_SHIFT:=1): 'SQ_TT_TOKEN_MASK_SHDEC_SHIFT', (SQ_TT_TOKEN_MASK_GFXUDEC_SHIFT:=2): 'SQ_TT_TOKEN_MASK_GFXUDEC_SHIFT', (SQ_TT_TOKEN_MASK_COMP_SHIFT:=3): 'SQ_TT_TOKEN_MASK_COMP_SHIFT', (SQ_TT_TOKEN_MASK_CONTEXT_SHIFT:=4): 'SQ_TT_TOKEN_MASK_CONTEXT_SHIFT', (SQ_TT_TOKEN_MASK_CONFIG_SHIFT:=5): 'SQ_TT_TOKEN_MASK_CONFIG_SHIFT', (SQ_TT_TOKEN_MASK_ALL_SHIFT:=6): 'SQ_TT_TOKEN_MASK_ALL_SHIFT', (SQ_TT_TOKEN_MASK_RSVD_SHIFT:=7): 'SQ_TT_TOKEN_MASK_RSVD_SHIFT'}
|
||||
SQ_TT_TOKEN_MASK_REG_INCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT:=0): 'SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT:=1): 'SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT:=2): 'SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT:=3): 'SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT:=4): 'SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT', (SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT:=5): 'SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT', (SQ_TT_TOKEN_EXCLUDE_REG_SHIFT:=6): 'SQ_TT_TOKEN_EXCLUDE_REG_SHIFT', (SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT:=7): 'SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT', (SQ_TT_TOKEN_EXCLUDE_INST_SHIFT:=8): 'SQ_TT_TOKEN_EXCLUDE_INST_SHIFT', (SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT:=9): 'SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT:=10): 'SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT:=11): 'SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT'}
|
||||
SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_UTIL_TIMER: dict[int, str] = {(SQ_TT_UTIL_TIMER_100_CLK:=0): 'SQ_TT_UTIL_TIMER_100_CLK', (SQ_TT_UTIL_TIMER_250_CLK:=1): 'SQ_TT_UTIL_TIMER_250_CLK'}
|
||||
SQ_TT_UTIL_TIMER: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_WAVESTART_MODE: dict[int, str] = {(SQ_TT_WAVESTART_MODE_SHORT:=0): 'SQ_TT_WAVESTART_MODE_SHORT', (SQ_TT_WAVESTART_MODE_ALLOC:=1): 'SQ_TT_WAVESTART_MODE_ALLOC', (SQ_TT_WAVESTART_MODE_PBB_ID:=2): 'SQ_TT_WAVESTART_MODE_PBB_ID'}
|
||||
SQ_TT_WAVESTART_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_WTYPE_INCLUDE: dict[int, str] = {(SQ_TT_WTYPE_INCLUDE_PS_BIT:=1): 'SQ_TT_WTYPE_INCLUDE_PS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD0_BIT:=2): 'SQ_TT_WTYPE_INCLUDE_RSVD0_BIT', (SQ_TT_WTYPE_INCLUDE_GS_BIT:=4): 'SQ_TT_WTYPE_INCLUDE_GS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD1_BIT:=8): 'SQ_TT_WTYPE_INCLUDE_RSVD1_BIT', (SQ_TT_WTYPE_INCLUDE_HS_BIT:=16): 'SQ_TT_WTYPE_INCLUDE_HS_BIT', (SQ_TT_WTYPE_INCLUDE_RSVD2_BIT:=32): 'SQ_TT_WTYPE_INCLUDE_RSVD2_BIT', (SQ_TT_WTYPE_INCLUDE_CS_BIT:=64): 'SQ_TT_WTYPE_INCLUDE_CS_BIT'}
|
||||
SQ_TT_WTYPE_INCLUDE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_WTYPE_INCLUDE_SHIFT: dict[int, str] = {(SQ_TT_WTYPE_INCLUDE_PS_SHIFT:=0): 'SQ_TT_WTYPE_INCLUDE_PS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD0_SHIFT:=1): 'SQ_TT_WTYPE_INCLUDE_RSVD0_SHIFT', (SQ_TT_WTYPE_INCLUDE_GS_SHIFT:=2): 'SQ_TT_WTYPE_INCLUDE_GS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD1_SHIFT:=3): 'SQ_TT_WTYPE_INCLUDE_RSVD1_SHIFT', (SQ_TT_WTYPE_INCLUDE_HS_SHIFT:=4): 'SQ_TT_WTYPE_INCLUDE_HS_SHIFT', (SQ_TT_WTYPE_INCLUDE_RSVD2_SHIFT:=5): 'SQ_TT_WTYPE_INCLUDE_RSVD2_SHIFT', (SQ_TT_WTYPE_INCLUDE_CS_SHIFT:=6): 'SQ_TT_WTYPE_INCLUDE_CS_SHIFT'}
|
||||
SQ_TT_WTYPE_INCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_CSCNTL_TYPE: dict[int, str] = {(CSCNTL_TYPE_TG:=0): 'CSCNTL_TYPE_TG', (CSCNTL_TYPE_STATE:=1): 'CSCNTL_TYPE_STATE', (CSCNTL_TYPE_EVENT:=2): 'CSCNTL_TYPE_EVENT', (CSCNTL_TYPE_PRIVATE:=3): 'CSCNTL_TYPE_PRIVATE'}
|
||||
CSCNTL_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
|
||||
CSDATA_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (Reserved_0x09:=9): 'Reserved_0x09', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (WAIT_SYNC:=21): 'WAIT_SYNC', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (FLUSH_ES_OUTPUT:=28): 'FLUSH_ES_OUTPUT', (BIN_CONF_OVERRIDE_CHECK:=29): 'BIN_CONF_OVERRIDE_CHECK', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_DRAW:=54): 'THREAD_TRACE_DRAW', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_LEGACY_PIPELINE:=62): 'ENABLE_LEGACY_PIPELINE', (DRAW_DONE:=63): 'DRAW_DONE'}
|
||||
VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32
|
||||
@@ -0,0 +1,24 @@
|
||||
# mypy: disable-error-code="empty-body"
|
||||
from __future__ import annotations
|
||||
import ctypes
|
||||
from typing import Literal, TypeAlias
|
||||
from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
|
||||
from tinygrad.runtime.support import c
|
||||
enum_MTYPE: dict[int, str] = {(MTYPE_C_RW_US:=0): 'MTYPE_C_RW_US', (MTYPE_RESERVED_1:=1): 'MTYPE_RESERVED_1', (MTYPE_C_RO_S:=2): 'MTYPE_C_RO_S', (MTYPE_UC:=3): 'MTYPE_UC', (MTYPE_C_RW_S:=4): 'MTYPE_C_RW_S', (MTYPE_RESERVED_5:=5): 'MTYPE_RESERVED_5', (MTYPE_C_RO_US:=6): 'MTYPE_C_RO_US', (MTYPE_RESERVED_7:=7): 'MTYPE_RESERVED_7'}
|
||||
MTYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_CSCNTL_TYPE: dict[int, str] = {(CSCNTL_TYPE_TG:=0): 'CSCNTL_TYPE_TG', (CSCNTL_TYPE_STATE:=1): 'CSCNTL_TYPE_STATE', (CSCNTL_TYPE_EVENT:=2): 'CSCNTL_TYPE_EVENT', (CSCNTL_TYPE_PRIVATE:=3): 'CSCNTL_TYPE_PRIVATE'}
|
||||
CSCNTL_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
|
||||
CSDATA_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (EVENT_STATE_CHANGE:=9): 'EVENT_STATE_CHANGE', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (WAIT_SYNC:=21): 'WAIT_SYNC', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (FLUSH_ES_OUTPUT:=28): 'FLUSH_ES_OUTPUT', (BIN_CONF_OVERRIDE_CHECK:=29): 'BIN_CONF_OVERRIDE_CHECK', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_DRAW:=54): 'THREAD_TRACE_DRAW', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_PIPELINE_NOT_USED:=62): 'ENABLE_PIPELINE_NOT_USED', (DRAW_DONE:=63): 'DRAW_DONE'}
|
||||
VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
|
||||
SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
|
||||
SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_REG_INCLUDE: dict[int, str] = {(SQ_TT_TOKEN_MASK_SQDEC_BIT:=1): 'SQ_TT_TOKEN_MASK_SQDEC_BIT', (SQ_TT_TOKEN_MASK_SHDEC_BIT:=2): 'SQ_TT_TOKEN_MASK_SHDEC_BIT', (SQ_TT_TOKEN_MASK_GFXUDEC_BIT:=4): 'SQ_TT_TOKEN_MASK_GFXUDEC_BIT', (SQ_TT_TOKEN_MASK_COMP_BIT:=8): 'SQ_TT_TOKEN_MASK_COMP_BIT', (SQ_TT_TOKEN_MASK_CONTEXT_BIT:=16): 'SQ_TT_TOKEN_MASK_CONTEXT_BIT'}
|
||||
SQ_TT_TOKEN_MASK_REG_INCLUDE: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: dict[int, str] = {(SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT:=0): 'SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT:=1): 'SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT', (SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT:=3): 'SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT'}
|
||||
SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT: TypeAlias = ctypes.c_uint32
|
||||
enum_SQ_TT_MODE: dict[int, str] = {(SQ_TT_MODE_OFF:=0): 'SQ_TT_MODE_OFF', (SQ_TT_MODE_ON:=1): 'SQ_TT_MODE_ON'}
|
||||
SQ_TT_MODE: TypeAlias = ctypes.c_uint32
|
||||
@@ -0,0 +1,16 @@
|
||||
# mypy: disable-error-code="empty-body"
|
||||
from __future__ import annotations
|
||||
import ctypes
|
||||
from typing import Literal, TypeAlias
|
||||
from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
|
||||
from tinygrad.runtime.support import c
|
||||
enum_MTYPE: dict[int, str] = {(MTYPE_NC:=0): 'MTYPE_NC', (MTYPE_WC:=1): 'MTYPE_WC', (MTYPE_RW:=1): 'MTYPE_RW', (MTYPE_CC:=2): 'MTYPE_CC', (MTYPE_UC:=3): 'MTYPE_UC'}
|
||||
MTYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ADDRESS_MODE: dict[int, str] = {(SH_MEM_ADDRESS_MODE_64:=0): 'SH_MEM_ADDRESS_MODE_64', (SH_MEM_ADDRESS_MODE_32:=1): 'SH_MEM_ADDRESS_MODE_32'}
|
||||
SH_MEM_ADDRESS_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_SH_MEM_ALIGNMENT_MODE: dict[int, str] = {(SH_MEM_ALIGNMENT_MODE_DWORD:=0): 'SH_MEM_ALIGNMENT_MODE_DWORD', (SH_MEM_ALIGNMENT_MODE_DWORD_STRICT:=1): 'SH_MEM_ALIGNMENT_MODE_DWORD_STRICT', (SH_MEM_ALIGNMENT_MODE_STRICT:=2): 'SH_MEM_ALIGNMENT_MODE_STRICT', (SH_MEM_ALIGNMENT_MODE_UNALIGNED:=3): 'SH_MEM_ALIGNMENT_MODE_UNALIGNED'}
|
||||
SH_MEM_ALIGNMENT_MODE: TypeAlias = ctypes.c_uint32
|
||||
enum_CSDATA_TYPE: dict[int, str] = {(CSDATA_TYPE_TG:=0): 'CSDATA_TYPE_TG', (CSDATA_TYPE_STATE:=1): 'CSDATA_TYPE_STATE', (CSDATA_TYPE_EVENT:=2): 'CSDATA_TYPE_EVENT', (CSDATA_TYPE_PRIVATE:=3): 'CSDATA_TYPE_PRIVATE'}
|
||||
CSDATA_TYPE: TypeAlias = ctypes.c_uint32
|
||||
enum_VGT_EVENT_TYPE: dict[int, str] = {(Reserved_0x00:=0): 'Reserved_0x00', (SAMPLE_STREAMOUTSTATS1:=1): 'SAMPLE_STREAMOUTSTATS1', (SAMPLE_STREAMOUTSTATS2:=2): 'SAMPLE_STREAMOUTSTATS2', (SAMPLE_STREAMOUTSTATS3:=3): 'SAMPLE_STREAMOUTSTATS3', (CACHE_FLUSH_TS:=4): 'CACHE_FLUSH_TS', (CONTEXT_DONE:=5): 'CONTEXT_DONE', (CACHE_FLUSH:=6): 'CACHE_FLUSH', (CS_PARTIAL_FLUSH:=7): 'CS_PARTIAL_FLUSH', (VGT_STREAMOUT_SYNC:=8): 'VGT_STREAMOUT_SYNC', (Reserved_0x09:=9): 'Reserved_0x09', (VGT_STREAMOUT_RESET:=10): 'VGT_STREAMOUT_RESET', (END_OF_PIPE_INCR_DE:=11): 'END_OF_PIPE_INCR_DE', (END_OF_PIPE_IB_END:=12): 'END_OF_PIPE_IB_END', (RST_PIX_CNT:=13): 'RST_PIX_CNT', (BREAK_BATCH:=14): 'BREAK_BATCH', (VS_PARTIAL_FLUSH:=15): 'VS_PARTIAL_FLUSH', (PS_PARTIAL_FLUSH:=16): 'PS_PARTIAL_FLUSH', (FLUSH_HS_OUTPUT:=17): 'FLUSH_HS_OUTPUT', (FLUSH_DFSM:=18): 'FLUSH_DFSM', (RESET_TO_LOWEST_VGT:=19): 'RESET_TO_LOWEST_VGT', (CACHE_FLUSH_AND_INV_TS_EVENT:=20): 'CACHE_FLUSH_AND_INV_TS_EVENT', (ZPASS_DONE:=21): 'ZPASS_DONE', (CACHE_FLUSH_AND_INV_EVENT:=22): 'CACHE_FLUSH_AND_INV_EVENT', (PERFCOUNTER_START:=23): 'PERFCOUNTER_START', (PERFCOUNTER_STOP:=24): 'PERFCOUNTER_STOP', (PIPELINESTAT_START:=25): 'PIPELINESTAT_START', (PIPELINESTAT_STOP:=26): 'PIPELINESTAT_STOP', (PERFCOUNTER_SAMPLE:=27): 'PERFCOUNTER_SAMPLE', (Available_0x1c:=28): 'Available_0x1c', (Available_0x1d:=29): 'Available_0x1d', (SAMPLE_PIPELINESTAT:=30): 'SAMPLE_PIPELINESTAT', (SO_VGTSTREAMOUT_FLUSH:=31): 'SO_VGTSTREAMOUT_FLUSH', (SAMPLE_STREAMOUTSTATS:=32): 'SAMPLE_STREAMOUTSTATS', (RESET_VTX_CNT:=33): 'RESET_VTX_CNT', (BLOCK_CONTEXT_DONE:=34): 'BLOCK_CONTEXT_DONE', (CS_CONTEXT_DONE:=35): 'CS_CONTEXT_DONE', (VGT_FLUSH:=36): 'VGT_FLUSH', (TGID_ROLLOVER:=37): 'TGID_ROLLOVER', (SQ_NON_EVENT:=38): 'SQ_NON_EVENT', (SC_SEND_DB_VPZ:=39): 'SC_SEND_DB_VPZ', (BOTTOM_OF_PIPE_TS:=40): 'BOTTOM_OF_PIPE_TS', (FLUSH_SX_TS:=41): 'FLUSH_SX_TS', (DB_CACHE_FLUSH_AND_INV:=42): 'DB_CACHE_FLUSH_AND_INV', (FLUSH_AND_INV_DB_DATA_TS:=43): 'FLUSH_AND_INV_DB_DATA_TS', (FLUSH_AND_INV_DB_META:=44): 'FLUSH_AND_INV_DB_META', (FLUSH_AND_INV_CB_DATA_TS:=45): 'FLUSH_AND_INV_CB_DATA_TS', (FLUSH_AND_INV_CB_META:=46): 'FLUSH_AND_INV_CB_META', (CS_DONE:=47): 'CS_DONE', (PS_DONE:=48): 'PS_DONE', (FLUSH_AND_INV_CB_PIXEL_DATA:=49): 'FLUSH_AND_INV_CB_PIXEL_DATA', (SX_CB_RAT_ACK_REQUEST:=50): 'SX_CB_RAT_ACK_REQUEST', (THREAD_TRACE_START:=51): 'THREAD_TRACE_START', (THREAD_TRACE_STOP:=52): 'THREAD_TRACE_STOP', (THREAD_TRACE_MARKER:=53): 'THREAD_TRACE_MARKER', (THREAD_TRACE_FLUSH:=54): 'THREAD_TRACE_FLUSH', (THREAD_TRACE_FINISH:=55): 'THREAD_TRACE_FINISH', (PIXEL_PIPE_STAT_CONTROL:=56): 'PIXEL_PIPE_STAT_CONTROL', (PIXEL_PIPE_STAT_DUMP:=57): 'PIXEL_PIPE_STAT_DUMP', (PIXEL_PIPE_STAT_RESET:=58): 'PIXEL_PIPE_STAT_RESET', (CONTEXT_SUSPEND:=59): 'CONTEXT_SUSPEND', (OFFCHIP_HS_DEALLOC:=60): 'OFFCHIP_HS_DEALLOC', (ENABLE_NGG_PIPELINE:=61): 'ENABLE_NGG_PIPELINE', (ENABLE_LEGACY_PIPELINE:=62): 'ENABLE_LEGACY_PIPELINE', (Reserved_0x3f:=63): 'Reserved_0x3f'}
|
||||
VGT_EVENT_TYPE: TypeAlias = ctypes.c_uint32
|
||||
@@ -33,7 +33,7 @@ class NullAllocator(Allocator['NullDevice']):
|
||||
class NullGraph(MultiGraphRunner):
|
||||
def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
|
||||
# description based on command, copied from HCQ graph
|
||||
if PROFILE: cpu_events.append(ProfileGraphEvent(ents:=[ProfileGraphEntry(self.device if runtime is not None else f"{self.device}:SDMA:0", \
|
||||
if PROFILE: cpu_events.append(ProfileGraphEvent(ents:=[ProfileGraphEntry(runtime.device if runtime is not None else f"{bufs[1].device}:SDMA:0", \
|
||||
runtime.name if runtime is not None else f"{bufs[1].device} -> {bufs[0].device}", i, i+1) \
|
||||
for i,((_,_,bufs,_),runtime) in enumerate(zip(self.calls, self.runtimes))], [], [perf_counter_us() for _ in range(len(ents)+1)]))
|
||||
return 1e-1
|
||||
|
||||
@@ -326,7 +326,7 @@ class AMDev:
|
||||
@functools.cached_property
|
||||
def hwid_names(self) -> dict[int, str]: return {v:k.removesuffix('_HWID') for k,v in vars(am).items() if k.endswith('_HWID') and isinstance(v, int)}
|
||||
|
||||
def _ip_module(self, prefix:str, hwip, prever_prefix:str=""): return import_module(prefix, self.ip_ver[hwip], prever_prefix)
|
||||
def _ip_module(self, prefix:str, hwip): return import_module(prefix, self.ip_ver[hwip])
|
||||
|
||||
def _build_regs(self):
|
||||
mods = [("mp", am.MP0_HWIP), ("hdp", am.HDP_HWIP), ("gc", am.GC_HWIP), ("mmhub", am.MMHUB_HWIP), ("osssys", am.OSSSYS_HWIP),
|
||||
|
||||
@@ -173,7 +173,7 @@ class AM_GMC(AM_IP):
|
||||
|
||||
class AM_SMU(AM_IP):
|
||||
def init_sw(self):
|
||||
self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP, prever_prefix='v')
|
||||
self.smu_mod = self.adev._ip_module("smu", am.MP1_HWIP)
|
||||
self.driver_table_paddr = self.adev.mm.palloc(0x4000, zero=False, boot=True)
|
||||
|
||||
def init_hw(self):
|
||||
|
||||
@@ -2,7 +2,6 @@ import functools, re, tinygrad.runtime.autogen.am
|
||||
from dataclasses import dataclass
|
||||
from tinygrad.helpers import getbits, fetch
|
||||
|
||||
AMDGPU_URL = "https://gitlab.com/linux-kernel/linux-next/-/raw/cf6d949a409e09539477d32dbe7c954e4852e744/drivers/gpu/drm/amd"
|
||||
ROCM_URL = "https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects"
|
||||
|
||||
@dataclass
|
||||
@@ -28,37 +27,18 @@ class AMDIP:
|
||||
if (name10:=name.replace('reg', 'mm')) in self.regs: return self.regs[name10]
|
||||
raise AttributeError(f"{self.name.upper()} has no register {name}")
|
||||
|
||||
def fixup_ip_version(ip:str, version:tuple[int, ...]) -> list[tuple[int, ...]]:
|
||||
# override versions
|
||||
def _apply_ovrd(ovrd:dict[tuple[int, ...], tuple[int, ...]]) -> tuple[int, ...]:
|
||||
for ver, ovrd_ver in ovrd.items():
|
||||
if version[:len(ver)] == ver: return ovrd_ver
|
||||
return version
|
||||
# load the greatest module with matching major version that's less than or equal to the target version
|
||||
# this is not universally correct, see below for an example, but appears reliable for recent gpus
|
||||
# https://github.com/torvalds/linux/blob/9207d47f966be9f4d52e7e0119ac2b7a7e366f3e/drivers/gpu/drm/amd/amdgpu/amdgpu_discovery.c#L3163
|
||||
def import_module(name:str, target:tuple[int, ...], submod=""):
|
||||
mod = getattr(tinygrad.runtime.autogen.am, submod) if submod else tinygrad.runtime.autogen.am
|
||||
if (children:=[c for c in mod.__all__ if c.startswith(name) and (v:=tuple(map(int, c.split('_')[1:])))[0] == target[0] and v <= target]):
|
||||
return getattr(mod, children[-1])
|
||||
raise ImportError(f"Failed to import {submod+'.' if submod else ''}{name} {'.'.join(map(str, target))}")
|
||||
|
||||
if ip in ['nbio', 'nbif']: version = _apply_ovrd({(7,3): (7,2,0)})
|
||||
elif ip in ['mp', 'smu']: version = _apply_ovrd({(14,0,3): (14,0,2)})
|
||||
elif ip in ['gc']: version = _apply_ovrd({(9,5,0): (9,4,3)})
|
||||
elif ip in ['sdma']: version = _apply_ovrd({(4,4,4): (4,4,2)})
|
||||
def header_download(file, url) -> str: return fetch(f"{url}/{file}", subdir="defines").read_text()
|
||||
|
||||
return [version, version[:2], version[:2]+(0,), version[:1]+(0, 0)]
|
||||
|
||||
def header_download(file, name=None, subdir="defines", url=AMDGPU_URL) -> str: return fetch(f"{url}/{file}", name=name, subdir=subdir).read_text()
|
||||
|
||||
def import_header(path:str, url=AMDGPU_URL):
|
||||
t = re.sub(r'//.*|/\*.*?\*/','', header_download(path, subdir="defines", url=url), flags=re.S)
|
||||
# TODO: refactor when clang2py is replaced
|
||||
return {k:int(v,0) for k,v in re.findall(r'\b([A-Za-z_]\w*)\s*=\s*(0x[0-9A-Fa-f]+|\d+)', t) + \
|
||||
re.findall(r'^\s*#\s*define\s+([A-Za-z_0-9]\w*)\s+(0x[0-9A-Fa-f]+|\d+)', t, re.M)}
|
||||
|
||||
def import_module(name:str, version:tuple[int, ...], version_prefix:str=""):
|
||||
for ver in fixup_ip_version(name, version):
|
||||
try: return getattr(tinygrad.runtime.autogen.am, f"{name}_{version_prefix}{'_'.join(map(str, ver))}")
|
||||
except AttributeError: pass
|
||||
raise ImportError(f"Failed to load autogen module for {name.upper()} {'.'.join(map(str, version))}")
|
||||
|
||||
def import_soc(ip):
|
||||
# rocm soc headers have more profiling enums than upstream linux
|
||||
return type("SOC", (object,), import_header(f"aqlprofile/linux/{({9: 'vega10', 10: 'navi10', 11: 'soc21', 12: 'soc24'}[ip[0]])}_enum.h", ROCM_URL))
|
||||
def import_soc(ip): return getattr(tinygrad.runtime.autogen.am, f"soc_{ip[0]}")
|
||||
|
||||
def import_pmc(ip) -> dict[str, tuple[str, int]]:
|
||||
res:dict[str, tuple[str, int]] = {}
|
||||
@@ -66,7 +46,7 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]:
|
||||
# NOTE: precise arch for mi300+, generic for others, since rocm headers lack some archs
|
||||
arch = f"gfx{ip[0]}{ip[1]:x}{ip[2]:x}" if ip[0] == 9 else f"gfx{ip[0]}"
|
||||
|
||||
for sec in header_download("rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml", url=ROCM_URL).split('- name: ')[1:]:
|
||||
for sec in header_download("rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml", ROCM_URL).split('- name: ')[1:]:
|
||||
for arch_spec in sec.split('- architectures:')[1:]:
|
||||
if arch in arch_spec and (block:=re.search(r'block:\s*([A-Za-z0-9_]+)', arch_spec)) and (ev:=re.search(r'event:\s*(\d+)', arch_spec)):
|
||||
res[sec.splitlines()[0].strip()] = (block.group(1), int(ev.group(1)))
|
||||
@@ -74,7 +54,4 @@ def import_pmc(ip) -> dict[str, tuple[str, int]]:
|
||||
return res
|
||||
|
||||
def import_asic_regs(prefix:str, version:tuple[int, ...], cls=AMDReg) -> dict[str, AMDReg]:
|
||||
from tinygrad.runtime.autogen.am import regs
|
||||
if (mods:=[m for m in regs.__all__ if m.startswith(prefix) and (v:=tuple(map(int, m.split('_')[1:])))[0] == version[0] and v <= version]):
|
||||
return {reg:cls(name=reg, offset=off, segment=seg, fields=fields) for reg,(off,seg,fields) in getattr(regs, mods[-1]).items()}
|
||||
raise ImportError(f"Failed to load ASIC registers for {prefix.upper()} {'.'.join(map(str, version))}")
|
||||
return {reg:cls(name=reg, offset=off, segment=seg, fields=fields) for reg,(off,seg,fields) in import_module(prefix, version, submod="regs").items()}
|
||||
|
||||
@@ -99,7 +99,8 @@ arc_families = ['alloc', 'copy', 'mutableCopy', 'new']
|
||||
|
||||
def normalize(a): return ("_" + n if keyword.iskeyword(n:=nm(a)) else n)
|
||||
|
||||
def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, macros=True, paths=[]):
|
||||
def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=False, errno=False, anon_names={}, types={}, macros=True, paths=[],
|
||||
patterns=[]):
|
||||
extras, lines, anoncnt, types, objc, fns = [], [], itertools.count().__next__, {k:(v,True) for k,v in types.items()}, False, set()
|
||||
|
||||
# ctypes automatically "unboxes" simple types
|
||||
@@ -227,6 +228,7 @@ def gen(name, files, dll="", args=[], prolog=[], rules=[], epilog=[], recsym=Fal
|
||||
while q:
|
||||
c = q.pop()
|
||||
if loc_file(loc(c)) != str(f) and (not recsym or c.kind not in (clang.CXCursor_FunctionDecl,)): continue
|
||||
if patterns and not any(re.match(p, nm(c)) for p in patterns): continue
|
||||
rollback = lines, types
|
||||
try:
|
||||
match c.kind:
|
||||
|
||||
@@ -290,8 +290,10 @@ def fast_idiv(target: Target, x: UOp, d: int, dont_cast=False) -> UOp|None:
|
||||
if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
|
||||
return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
|
||||
# before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
|
||||
# use explicit Ops.IDIV (trunc) since the recursion assumes trunc semantics throughout
|
||||
if (largest_factor_of_two_in_d := (d & -d)) > 1:
|
||||
if (ret:=fast_idiv(target, x//largest_factor_of_two_in_d, d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
|
||||
if (ret:=fast_idiv(target, x.alu(Ops.IDIV, x.const_like(largest_factor_of_two_in_d)),
|
||||
d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
|
||||
if dont_cast: return None
|
||||
# promo_lattice needs to return an unsigned type if the type is unsigned
|
||||
if dtypes.is_int(next_dtype := promo_lattice[x.dtype.scalar()][-1]) and is_dtype_supported(next_dtype, target):
|
||||
@@ -440,11 +442,11 @@ def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool)
|
||||
if Ops.SQRT not in ops or force_transcendental: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
|
||||
return PatternMatcher(pat)
|
||||
|
||||
def floordiv_to_idiv(d:UOp, a:UOp, b:UOp) -> UOp:
|
||||
def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
|
||||
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.IDIV, b)
|
||||
return a.alu(Ops.IDIV, b) - (a.alu(Ops.MOD, b).ne(0) & (a<0).ne(b<0)).cast(d.dtype)
|
||||
return a.alu(Ops.IDIV, b) - (a.alu(Ops.MOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
|
||||
|
||||
def floormod_to_mod(d:UOp, a:UOp, b:UOp) -> UOp:
|
||||
def floormod_to_mod(a:UOp, b:UOp) -> UOp:
|
||||
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.MOD, b)
|
||||
r = a.alu(Ops.MOD, b)
|
||||
# use where instead of mul to avoid being fused into MULACC (which int64 long-decomp doesn't handle)
|
||||
@@ -453,30 +455,34 @@ def floormod_to_mod(d:UOp, a:UOp, b:UOp) -> UOp:
|
||||
powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
|
||||
@functools.cache
|
||||
def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> PatternMatcher:
|
||||
pat: list[tuple[UPat, Callable]] = [
|
||||
(UPat(Ops.FLOORDIV, name="d", src=(UPat.var("a"), UPat.var("b"))), floordiv_to_idiv),
|
||||
(UPat(Ops.FLOORMOD, name="d", src=(UPat.var("a"), UPat.var("b"))), floormod_to_mod),
|
||||
]
|
||||
pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
|
||||
# FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
|
||||
if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None))
|
||||
pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
|
||||
# no real hardware supports THREEFRY, but NullRenderer does
|
||||
if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32))
|
||||
# MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends)
|
||||
if Ops.MAX not in ops and Ops.CMPLT in ops: pat.append((UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])))
|
||||
# rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
|
||||
# TODO: drop the x.vmin>=0 guard once UOp `%` lowers to FLOORMOD instead of MOD
|
||||
if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"),
|
||||
lambda x,c: x & (c.arg-1) if c.arg in powers_of_two and x.vmin >= 0 else None)]
|
||||
if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(),
|
||||
lambda x,y: (x | y).logical_not())]
|
||||
# rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
|
||||
if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
if Ops.SHR in ops:
|
||||
# no reason to check x<0 for uints
|
||||
pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
|
||||
c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] # (x+(x<0).where(c-1, 0)) >> v
|
||||
# uint IDIV by 2**v -> x >> v (FLOORDIV is lowered to IDIV by the rule above before reaching here)
|
||||
pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))),
|
||||
lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
# signed IDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v
|
||||
pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))),
|
||||
lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
|
||||
if (v:=powers_of_two.get(c.arg, 0)) else None)]
|
||||
if not disable_fast_idiv:
|
||||
pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d", vec=False), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
|
||||
pat += [(UPat.var("x", dtypes.ints)%UPat.var("d"), lambda x, d: x-d*(x//d))]
|
||||
# fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
|
||||
pat += [(UPat(Ops.IDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d", vec=False))),
|
||||
lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
|
||||
# rewrite raw MOD -> x - d*IDIV(x,d) so fast_idiv can pick up the IDIV. only on non-negative inputs;
|
||||
# avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.MOD as an implementation detail)
|
||||
pat += [(UPat(Ops.MOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
|
||||
lambda x, d: x - d * x.alu(Ops.IDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
|
||||
if Ops.NEG in ops:
|
||||
pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
|
||||
if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]
|
||||
|
||||
+46
-53
@@ -1,19 +1,19 @@
|
||||
import functools, itertools, math
|
||||
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
|
||||
from tinygrad.helpers import floordiv, floormod, unwrap
|
||||
|
||||
# NOTE: this cache is only on index UOps
|
||||
@functools.cache
|
||||
def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
|
||||
def fold_divmod_general(d: UOp) -> UOp|None:
|
||||
x, y = d.src
|
||||
|
||||
# cancel_divmod: simple cancel div/mod case when the range of the numerator lies within a single denominator interval
|
||||
x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
|
||||
assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
|
||||
if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
|
||||
if y_min*y_max > 0 and (qv:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
|
||||
return x - qv*y if d.op is Ops.MOD else d.const_like(qv)
|
||||
if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
|
||||
if y_min*y_max > 0 and (qv:=floordiv(x_min,y_min)) == floordiv(x_min,y_max) == floordiv(x_max,y_min) == floordiv(x_max,y_max):
|
||||
return x - qv*y if d.op is Ops.FLOORMOD else d.const_like(qv)
|
||||
|
||||
# split uops for the rest of the processing
|
||||
x_peeled, const = x.pop_const()
|
||||
@@ -22,19 +22,20 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
|
||||
# ** Constant Denominator Rules **
|
||||
# these rules strictly require y to be a scalar constant > 0
|
||||
if y.op is Ops.CONST and (c := y.arg) > 0:
|
||||
# nested_div_mod: (x%(k*c))//c -> (x//c)%k, and (x%(k*c))%c -> x%c
|
||||
if x.op is Ops.MOD and (k := x.src[1].divides(c)) is not None:
|
||||
return x.src[0] // y % k if d.op is Ops.IDIV else x.src[0] % y
|
||||
# nested_div_mod: (x%(k*c))//c -> (x//c)%k (requires k>0), and (x%(k*c))%c -> x%c
|
||||
if x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None:
|
||||
if d.op is Ops.FLOORMOD: return x.src[0] % y
|
||||
if k > 0: return x.src[0] // y % k
|
||||
|
||||
# remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2, requires non-negative sums
|
||||
if d.op is Ops.MOD and x.vmin >= 0:
|
||||
# remove_nested_mod in sum: (a%4 + b)%2 -> (a+b)%2
|
||||
if d.op is Ops.FLOORMOD:
|
||||
new_xs, changed = [], False
|
||||
for u in uops_no_const:
|
||||
if u.op is Ops.MOD and u.src[1].divides(c) is not None:
|
||||
if u.op is Ops.FLOORMOD and u.src[1].divides(c) is not None:
|
||||
u = u.src[0]
|
||||
changed = True
|
||||
new_xs.append(u)
|
||||
if changed and (new_x:=(UOp.usum(*new_xs) + const)).vmin >= 0: return new_x % y
|
||||
if changed: return (UOp.usum(*new_xs) + const) % y
|
||||
|
||||
# Shared decomposition for folding rules
|
||||
decomp = [(u.divides(f:=u.const_factor()),f) for u in uops_no_const]
|
||||
@@ -42,40 +43,39 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
|
||||
|
||||
# fold_binary_numerator: fold if expression has one non-constant term that takes on two values
|
||||
if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
|
||||
y1 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmin+const, c)
|
||||
y2 = (cmod if d.op is Ops.MOD else cdiv)(factors[0]*v.vmax+const, c)
|
||||
y1 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmin+const, c)
|
||||
y2 = (floormod if d.op is Ops.FLOORMOD else floordiv)(factors[0]*v.vmax+const, c)
|
||||
return (y2-y1)*(v-v.vmin) + y1
|
||||
|
||||
# fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
|
||||
if not (x.vmin<0 and correct_divmod_folding):
|
||||
# when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
|
||||
rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
|
||||
for rems in itertools.product(*rem_choices):
|
||||
if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
|
||||
if d.op is Ops.MOD: return rem - rem.vmin//c*c
|
||||
return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
|
||||
# when f%c == c//2, abs(r) == abs(r-c) is a tie, try both signs since either may fit in one period
|
||||
rem_choices = [(r, r-c) if (r:=f%c)*2 == c else (min(r, r-c, key=abs),) for f in factors]
|
||||
for rems in itertools.product(*rem_choices):
|
||||
if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
|
||||
if d.op is Ops.FLOORMOD: return rem - rem.vmin//c*c
|
||||
return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + const//c + rem.vmin//c
|
||||
|
||||
# gcd_with_remainder: factor out common gcd from numerator
|
||||
if x.vmin >= 0 and (g:=math.gcd(*factors, c)) > 1:
|
||||
if (g:=math.gcd(*factors, c)) > 1:
|
||||
new_x = unwrap(x_peeled.divides(g)).simplify() + (const//g)%(c//g)
|
||||
if new_x.vmin >= 0:
|
||||
if d.op is Ops.MOD: return new_x % (c//g) * g + const%g
|
||||
if d.op is Ops.FLOORMOD: return new_x % (c//g) * g + const%g
|
||||
return new_x // (c//g) + const//c
|
||||
|
||||
# nest_by_factor: x//c -> (x//f)//(c//f), x%c -> (x//f%(c//f))*f + b where b=x%f
|
||||
if x.vmin >= 0:
|
||||
results = []
|
||||
for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}:
|
||||
if (newxs := fold_divmod_general(x//div, correct_divmod_folding)) is not None and newxs.vmin >= 0:
|
||||
if d.op is Ops.IDIV:
|
||||
results.append((len(newxs.backward_slice), newxs // (c // div)))
|
||||
else:
|
||||
b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
|
||||
if const % div: b_parts.append(x.const_like(const % div))
|
||||
b = UOp.usum(*b_parts) if b_parts else x.const_like(0)
|
||||
if 0 <= b.vmin and b.vmax < div:
|
||||
results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
|
||||
if results: return min(results, key=lambda r: r[0])[1]
|
||||
# FLOORDIV identity holds for any sign of x; FLOORMOD reconstruction needs x.vmin>=0
|
||||
results = []
|
||||
for div in {abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and 1 < abs(f) < c and (c%f)==0}:
|
||||
if (newxs := fold_divmod_general(x//div)) is not None:
|
||||
if d.op is Ops.FLOORDIV:
|
||||
results.append((len(newxs.backward_slice), newxs // (c // div)))
|
||||
elif x.vmin >= 0 and newxs.vmin >= 0:
|
||||
b_parts = [f%div*t for f, t in zip(factors, terms) if f%div]
|
||||
if const % div: b_parts.append(x.const_like(const % div))
|
||||
b = UOp.usum(*b_parts) if b_parts else x.const_like(0)
|
||||
if 0 <= b.vmin and b.vmax < div:
|
||||
results.append((len((r:=(newxs % x.ufix(c//div))*div + b).backward_slice), r))
|
||||
if results: return min(results, key=lambda r: r[0])[1]
|
||||
|
||||
# ** Variable Denominator / Fallback Rules **
|
||||
# These rules apply to variables OR constants that failed the checks above.
|
||||
@@ -86,7 +86,7 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
|
||||
gcd = UOp.gcd(*all_uops, y).simplify()
|
||||
if not (gcd.op is Ops.CONST and gcd.arg==1):
|
||||
ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd)))
|
||||
return ret*gcd if d.op is Ops.MOD else ret
|
||||
return ret*gcd if d.op is Ops.FLOORMOD else ret
|
||||
|
||||
# factor_remainder: (d*x+y)//d -> x+y//d
|
||||
if y.vmin<0 or x.vmin<0: return None
|
||||
@@ -95,29 +95,22 @@ def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
|
||||
if (q:=u.divide_exact(y)) is not None: quo.append(q)
|
||||
elif y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c:
|
||||
rem.append(u.divides(c)*(c%y.arg))
|
||||
quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.IDIV else u.const_like(0))
|
||||
quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0))
|
||||
else: rem.append(u)
|
||||
|
||||
if not quo: return None
|
||||
new_x = sum(rem)+x.const_like(0)
|
||||
if new_x.vmin<0: return None
|
||||
return new_x%y if d.op is Ops.MOD else new_x//y+sum(quo)
|
||||
return new_x%y if d.op is Ops.FLOORMOD else new_x//y+sum(quo)
|
||||
|
||||
div_and_mod_symbolic = PatternMatcher([
|
||||
# ** 1. Fast Inline Rules **
|
||||
((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
|
||||
if c.vmin>0 and d.vmin>0 and x.vmin>=0 and a.vmin>=0 else None), # (x//c+a)//d -> (x+a*c)//(c*d)
|
||||
(UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
|
||||
(UPat.var("x", dtypes.weakint) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <= 0 else None),
|
||||
((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
|
||||
lambda x,c,n,d: ((x+c.arg%d.arg)//d + c.arg//d.arg) if c.arg%d.arg!=c.arg and x.vmin>=0 and n.vmin>=0 and d.arg>0 else None),
|
||||
((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
|
||||
lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
|
||||
# (x//c+a)//d -> (x+a*c)//(c*d) for c>0, d>0
|
||||
((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if c.vmin>0 and d.vmin>0 else None),
|
||||
# (x+c)//d -> (x+c%d)//d + c//d for d>0 (split out the multiple of d in the constant)
|
||||
((UPat.var("x", dtypes.weakint)+UPat.cvar("c", vec=False))//UPat.cvar("d", vec=False),
|
||||
lambda x,c,d: (x+c.arg%d.arg)//d + c.arg//d.arg if c.arg%d.arg!=c.arg and d.arg>0 else None),
|
||||
|
||||
# ** 2. Slow Rules **
|
||||
(UPat((Ops.IDIV, Ops.MOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d, bool(CORRECT_DIVMOD_FOLDING))),
|
||||
|
||||
# NOTE: these have to go at the bottom or TestSymbolicOps.test_var loops
|
||||
(UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
|
||||
(UPat.var("x", dtypes.weakint) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
|
||||
])
|
||||
(UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)),
|
||||
])
|
||||
|
||||
+3
-1
@@ -870,9 +870,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
return min(vals:=(cdiv(s0_vmin, s1_vmin), cdiv(s0_vmin, s1_vmax), cdiv(s0_vmax, s1_vmin), cdiv(s0_vmax, s1_vmax))), max(vals)
|
||||
if self.op is Ops.FLOORDIV:
|
||||
assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
|
||||
if s0_vmin > s0_vmax: return 0, 0 # numerator range is empty (e.g. RANGE with end=0)
|
||||
if s1_vmin*s1_vmax>0: return min(vals:=(s0_vmin//s1_vmin, s0_vmin//s1_vmax, s0_vmax//s1_vmin, s0_vmax//s1_vmax)), max(vals)
|
||||
if self.op is Ops.FLOORMOD:
|
||||
assert isinstance(s0_vmin, int) and isinstance(s0_vmax, int) and isinstance(s1_vmin, int) and isinstance(s1_vmax, int)
|
||||
if s0_vmin > s0_vmax: return 0, 0 # numerator range is empty (e.g. RANGE with end=0)
|
||||
if (c:=s1_vmin) == s1_vmax > 0: return (s0_vmin%c, s0_vmax%c) if s0_vmin//c == s0_vmax//c else (0, c-1)
|
||||
if (c:=s1_vmin) == s1_vmax < 0: return (s0_vmin%c, s0_vmax%c) if s0_vmin//c == s0_vmax//c else (c+1, 0)
|
||||
if s1_vmin > 0: return (0, s1_vmax-1)
|
||||
@@ -907,7 +909,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
|
||||
# TODO: sanitize varnames, or don't use naked eval while staying fast
|
||||
ret = _render_with_splits(list(sself.toposort()), renderer_infer, {sself})
|
||||
lines = [f" {k}={v}" for k,v in ret.items() if k != "ast"] + [f" return {ret['ast']}"]
|
||||
ns: dict[str, Any] = {"max": max, "cdiv": cdiv, "cmod": cmod, "bitcast": bitcast, "dtypes": dtypes}
|
||||
ns: dict[str, Any] = {"max": max, "cdiv": cdiv, "cmod": cmod, "floordiv": floordiv, "floormod": floormod, "bitcast": bitcast, "dtypes": dtypes}
|
||||
exec(f"def _f({','.join(varnames)}):\n"+'\n'.join(lines), ns) # pylint: disable=exec-used
|
||||
return ns["_f"], varnames
|
||||
|
||||
|
||||
+13
-6
@@ -23,10 +23,10 @@ def print_uops(uops:list[UOp]):
|
||||
print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}")
|
||||
|
||||
# for debug
|
||||
syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.IDIV: "//", Ops.MOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
|
||||
syms = { Ops.ADD: "+", Ops.SUB: "-", Ops.FLOORDIV: "//", Ops.FLOORMOD: "%", Ops.SHL: "<<", Ops.SHR: ">>",
|
||||
Ops.MUL: "*", Ops.CMPLT: "<", Ops.CMPNE: "!=", Ops.AND: "&", Ops.OR: "|", Ops.XOR: "^"}
|
||||
# comparison operators are not in here because they are chained in python, not left-associative
|
||||
precedence = {Ops.MUL:1, Ops.IDIV:1, Ops.MOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
|
||||
precedence = {Ops.MUL:1, Ops.FLOORDIV:1, Ops.FLOORMOD:1, Ops.ADD:2, Ops.SUB:2, Ops.SHL:3, Ops.SHR:3, Ops.AND:4, Ops.XOR:5, Ops.OR:6}
|
||||
def strip_binary_parens(x:UOp, left:str, right:str, code_for_op) -> str:
|
||||
if x.op not in precedence: return code_for_op(left, right)
|
||||
return code_for_op(strip_parens(left) if precedence.get(x.src[0].op,99)<=precedence[x.op] else left, strip_parens(right) if
|
||||
@@ -46,6 +46,8 @@ renderer = PatternMatcher([
|
||||
(UPat(Ops.MAX, name="x"), lambda ctx,x: f"max({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.MULACC, name="x"), lambda ctx,x: f"({ctx[x.src[0]]}*{ctx[x.src[1]]}+{ctx[x.src[2]]})"),
|
||||
(UPat(Ops.WHERE, name="x"), lambda ctx,x: f"({ctx[x.src[1]]} if {ctx[x.src[0]]} else {ctx[x.src[2]]})"),
|
||||
(UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(set(syms.keys()), name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat((Ops.INDEX, Ops.BUFFERIZE), name="x"), lambda x, ctx: ''.join([f"[{strip_parens(ctx[y])}]" for y in x.src[1:]])),
|
||||
(UPat(Ops.STACK, name="x"),
|
||||
@@ -56,6 +58,8 @@ renderer = PatternMatcher([
|
||||
renderer_infer = PatternMatcher([
|
||||
(UPat(Ops.MOD, name="x"), lambda ctx,x: f"cmod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.IDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"),
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"),
|
||||
]) + renderer
|
||||
|
||||
@@ -99,13 +103,16 @@ pm_pyrender_extra = PatternMatcher([
|
||||
# TODO: movement ops simplify stuff, this can break SPEC=2
|
||||
#(UPat(GroupOp.Movement, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({render_marg(ctx,x)})"),
|
||||
# NOTE: CMPNE doesn't work cause there's no __rne__
|
||||
# explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render IDIV/MOD via their named methods
|
||||
(UPat(Ops.IDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.idiv({ctx[x.src[1]]})"),
|
||||
(UPat(Ops.MOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.MOD, {ctx[x.src[1]]})"),
|
||||
# NOTE: only match CONSTs without UNIQUE (len(src)==1), unique_const needs explicit rendering
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.CMPNE, Ops.IDIV, Ops.MOD}, src=(UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="y"), UPat(name="z")), name="x"),
|
||||
lambda ctx,x,y,z: strip_binary_parens(x, str(y.arg), ctx[z], lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
# NOTE: sub doesn't work cause it's written as add/mul
|
||||
(UPat(set(syms.keys())-{Ops.SUB}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"), lambda ctx,x,y,z:
|
||||
strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat(set(syms.keys())-{Ops.SUB}, name="x"), lambda ctx,x:
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.IDIV, Ops.MOD}, src=(UPat(name="y"), UPat(Ops.CONST, src=(UPat(Ops.DEVICE),), name="z")), name="x"),
|
||||
lambda ctx,x,y,z: strip_binary_parens(x, ctx[y], str(z.arg), lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat(set(syms.keys())-{Ops.SUB, Ops.IDIV, Ops.MOD}, name="x"), lambda ctx,x:
|
||||
strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")),
|
||||
(UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"),
|
||||
(UPat(sugar, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}("+', '.join([ctx[y] for y in x.src[1:]] + \
|
||||
|
||||
+12
-11
@@ -28,8 +28,8 @@ invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
|
||||
def fold_add_divmod_recombine(x:UOp) -> UOp|None:
|
||||
terms = list(x.split_uop(Ops.ADD))
|
||||
for i,u in enumerate(terms):
|
||||
if u.op is Ops.MOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
|
||||
elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.MOD and m.src[1].op is Ops.CONST:
|
||||
if u.op is Ops.FLOORMOD and u.src[1].op is Ops.CONST: base, div, mul = u.src[0], u.src[1].arg, 1
|
||||
elif u.op is Ops.MUL and u.src[1].op is Ops.CONST and (m:=u.src[0]).op is Ops.FLOORMOD and m.src[1].op is Ops.CONST:
|
||||
base, div, mul = m.src[0], m.src[1].arg, u.src[1].arg
|
||||
else: continue
|
||||
for j,v in enumerate(terms):
|
||||
@@ -37,13 +37,13 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
|
||||
if v.op is not Ops.MUL or v.src[1].op is not Ops.CONST or v.src[1].arg != div*mul: continue
|
||||
q, exact = v.src[0], False
|
||||
# (base%div)*mul + (base//div)*(div*mul) -> base*mul
|
||||
if q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
|
||||
if q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[1].arg == div: exact = q.src[0] is base
|
||||
# ((base//d)%div)*mul + (base//(d*div))*(div*mul) -> (base//d)*mul
|
||||
if not exact and base.op is Ops.IDIV and base.src[1].op is Ops.CONST:
|
||||
exact = q.op is Ops.IDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
|
||||
if not exact and base.op is Ops.FLOORDIV and base.src[1].op is Ops.CONST:
|
||||
exact = q.op is Ops.FLOORDIV and q.src[1].op is Ops.CONST and q.src[0] is base.src[0] and q.src[1].arg == base.src[1].arg*div
|
||||
if exact: return (base*mul).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
|
||||
# ((base//div)%d)*div + base%div -> base%(div*d)
|
||||
if mul == 1 and div > 0 and q.op is Ops.MOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.IDIV:
|
||||
if mul == 1 and div > 0 and q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and q.src[0].op is Ops.FLOORDIV:
|
||||
if q.src[0].src[0] is base and q.src[0].src[1].op is Ops.CONST and q.src[0].src[1].arg == div:
|
||||
return (base % (div*d)).usum(*[t for k,t in enumerate(terms) if k not in (i,j)])
|
||||
return None
|
||||
@@ -244,7 +244,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
((UPat.var("y")+UPat.var("c").where(UPat.var("t"), UPat.var("f"))) + UPat.var("c").where(UPat.var("tt"), UPat.var("ff")), \
|
||||
lambda y,c,t,tt,f,ff: y+c.where(t+tt, f+ff) if t.op == tt.op == Ops.CONST or f.op == ff.op == Ops.CONST else None),
|
||||
# ALU/variable min==max -> CONST
|
||||
(UPat({Ops.CMPLT, Ops.CMPNE, Ops.IDIV, Ops.MOD, Ops.DEFINE_VAR, Ops.BIND, Ops.SPECIAL}, name="x"),
|
||||
(UPat({Ops.CMPLT, Ops.CMPNE, Ops.FLOORDIV, Ops.FLOORMOD, Ops.DEFINE_VAR, Ops.BIND, Ops.SPECIAL}, name="x"),
|
||||
lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
(UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
|
||||
# max folding
|
||||
@@ -255,7 +255,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
*((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
|
||||
lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2))) for op in GroupOp.Associative),
|
||||
((UPat.cvar("c0") + UPat.var("x")) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)), # c0 + x < c1 -> x < c1 - c0
|
||||
((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2)), # (x//c1)//c2 -> x//(c1*c2)
|
||||
# (x//c1)//c2 -> x//(c1*c2) for c2>0
|
||||
((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
|
||||
# ** lt **
|
||||
# c0*x<c1 for positive int c0,c1
|
||||
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
|
||||
@@ -263,9 +264,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
|
||||
# c0*x<c1 for negative int c0 and non-positive c1
|
||||
((UPat.cvar("c0", vec=False)*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1", vec=False),
|
||||
lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
|
||||
# x//d<c
|
||||
# x//d<c -> x<c*d for d>0
|
||||
((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d", vec=False))<UPat.cvar("c", vec=False),
|
||||
lambda x,d,c: (x<(c.arg*d.arg) if c.arg > 0 else x<(c.arg*d.arg-(d.arg-1))) if d.arg > 0 else None),
|
||||
lambda x,d,c: x<(c.arg*d.arg) if d.arg > 0 else None),
|
||||
# ** move add/mul consts to end (NOTE: this is still happening before constant folding) **
|
||||
((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1),
|
||||
((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1),
|
||||
@@ -408,7 +409,7 @@ pm_move_where_on_load = PatternMatcher([
|
||||
def gated_given_valid(cond:UOp, x:UOp, i:UOp) -> UOp|None:
|
||||
if x.dtype.scalar() is not dtypes.weakint: return None
|
||||
# Skip if x contains DIV/MOD AND IMAGE mode is enabled -> image index e.g. openpilot
|
||||
if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.IDIV, Ops.MOD): return None
|
||||
if IMAGE.value > 0 and x.op_in_backward_slice_with_self(Ops.IDIV, Ops.MOD, Ops.FLOORDIV, Ops.FLOORMOD): return None
|
||||
return cond.where(uop_given_valid(cond, x, try_simplex=False), i)
|
||||
|
||||
# TODO: this is O(number of WHERE * number of node)
|
||||
|
||||
Reference in New Issue
Block a user