mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 22:06:06 +00:00
work
This commit is contained in:
@@ -105,7 +105,8 @@ def disasm(inst: Inst) -> str:
|
||||
if op_name == 'v_nop': return 'v_nop'
|
||||
if op_name == 'v_pipeflush': return 'v_pipeflush'
|
||||
parts = op_name.split('_')
|
||||
is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
|
||||
# cvt instructions use full 32-bit registers (f16 result in bits[15:0]), not packed halves
|
||||
is_16bit_dst = 'cvt' not in op_name and (any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES))
|
||||
is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name and 'cvt' not in op_name
|
||||
_F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
|
||||
is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
|
||||
|
||||
@@ -1,802 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Tests for the RDNA assembly renderer"""
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, dtypes
|
||||
from tinygrad.helpers import getenv
|
||||
|
||||
# Skip tests if not on AMD RDNA device
|
||||
AMD_RDNA = getenv("AMD_RDNA", 0)
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNABasic(unittest.TestCase):
|
||||
"""Basic functionality tests"""
|
||||
|
||||
def test_basic_half_load(self):
|
||||
"""Test basic half-precision load and sum"""
|
||||
a = Tensor(np.arange(16, dtype=np.float16).reshape(1, 16))
|
||||
b = a.sum()
|
||||
b.realize()
|
||||
np.testing.assert_allclose(b.numpy(), np.arange(16).sum(), rtol=1e-2)
|
||||
|
||||
def test_basic_float_load(self):
|
||||
"""Test basic float load and sum"""
|
||||
a = Tensor(np.arange(16, dtype=np.float32).reshape(1, 16))
|
||||
b = a.sum()
|
||||
b.realize()
|
||||
np.testing.assert_allclose(b.numpy(), np.arange(16).sum(), rtol=1e-5)
|
||||
|
||||
def test_elementwise_add(self):
|
||||
"""Test elementwise addition"""
|
||||
a = Tensor([1.0, 2.0, 3.0, 4.0])
|
||||
b = Tensor([5.0, 6.0, 7.0, 8.0])
|
||||
c = a + b
|
||||
c.realize()
|
||||
np.testing.assert_allclose(c.numpy(), [6.0, 8.0, 10.0, 12.0])
|
||||
|
||||
def test_elementwise_mul(self):
|
||||
"""Test elementwise multiplication"""
|
||||
a = Tensor([1.0, 2.0, 3.0, 4.0])
|
||||
b = Tensor([2.0, 2.0, 2.0, 2.0])
|
||||
c = a * b
|
||||
c.realize()
|
||||
np.testing.assert_allclose(c.numpy(), [2.0, 4.0, 6.0, 8.0])
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNAGatedLoads(unittest.TestCase):
|
||||
"""Tests for gated loads (WHERE+LOAD patterns)"""
|
||||
|
||||
def test_relu_backward_gated_load(self):
|
||||
"""Test that relu backward correctly gates memory access"""
|
||||
Tensor.training = True
|
||||
x = Tensor(np.array([[-1, 2], [3, -4]], dtype=np.float32), device="AMD", requires_grad=True)
|
||||
y = x.relu()
|
||||
loss = y.sum()
|
||||
loss.backward()
|
||||
# Gradient should be 1 where x > 0, 0 elsewhere
|
||||
expected = np.array([[0, 1], [1, 0]], dtype=np.float32)
|
||||
np.testing.assert_allclose(x.grad.numpy(), expected)
|
||||
|
||||
def test_max_pool2d_backward_gated_load(self):
|
||||
"""Test that max_pool2d backward correctly gates memory access"""
|
||||
Tensor.training = True
|
||||
x = Tensor.rand(2, 4, 8, 8, device="AMD", requires_grad=True)
|
||||
y = x.max_pool2d(2)
|
||||
loss = y.sum()
|
||||
loss.backward()
|
||||
self.assertIsNotNone(x.grad)
|
||||
# Just verify it runs without GPU fault
|
||||
|
||||
def test_sparse_categorical_crossentropy(self):
|
||||
"""Test sparse_categorical_crossentropy with index-based gating"""
|
||||
Tensor.training = True
|
||||
logits = Tensor.rand(4, 10, device="AMD", requires_grad=True)
|
||||
targets = Tensor([0, 3, 5, 9], device="AMD")
|
||||
loss = logits.sparse_categorical_crossentropy(targets)
|
||||
loss.backward()
|
||||
self.assertIsNotNone(logits.grad)
|
||||
# Gradients for each row should sum to 0 (since exp/sum normalizes)
|
||||
grad_sum = logits.grad.numpy().sum(axis=1)
|
||||
np.testing.assert_allclose(grad_sum, np.zeros(4), atol=1e-5)
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNA64BitOperations(unittest.TestCase):
|
||||
"""Tests for 64-bit integer operations (used in division-by-multiplication pattern)"""
|
||||
|
||||
@unittest.skip("Division by constant requires fix for quotient register preservation in RDNA renderer")
|
||||
def test_integer_division_optimization(self):
|
||||
"""Test that integer division uses 64-bit MUL+SHR pattern correctly"""
|
||||
# Integer division by constants uses a mul+shift pattern that requires 64-bit ops
|
||||
a = Tensor([10, 20, 30, 40], dtype=dtypes.int32, device="AMD")
|
||||
b = a // 7
|
||||
expected = np.array([1, 2, 4, 5], dtype=np.int32)
|
||||
np.testing.assert_array_equal(b.numpy(), expected)
|
||||
|
||||
@unittest.skip("Modulo by constant requires fix for quotient register preservation in RDNA renderer")
|
||||
def test_modulo_operation(self):
|
||||
"""Test that modulo uses 64-bit intermediate operations correctly"""
|
||||
a = Tensor([10, 20, 30, 40], dtype=dtypes.int32, device="AMD")
|
||||
b = a % 7
|
||||
expected = np.array([3, 6, 2, 5], dtype=np.int32)
|
||||
np.testing.assert_array_equal(b.numpy(), expected)
|
||||
|
||||
def test_integer_division_by_tensor(self):
|
||||
"""Test that integer division by tensor works correctly (uses float conversion)"""
|
||||
a = Tensor([10, 20, 30, 40], dtype=dtypes.int32, device="AMD")
|
||||
b = Tensor([7, 7, 7, 7], dtype=dtypes.int32, device="AMD")
|
||||
c = a // b
|
||||
expected = np.array([1, 2, 4, 5], dtype=np.int32)
|
||||
np.testing.assert_array_equal(c.numpy(), expected)
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNATraining(unittest.TestCase):
|
||||
"""Tests for training/backward pass functionality"""
|
||||
|
||||
def test_simple_mlp_training(self):
|
||||
"""Test a simple 2-layer MLP training step"""
|
||||
Tensor.training = True
|
||||
np.random.seed(42)
|
||||
|
||||
w1 = Tensor.scaled_uniform(16, 8, requires_grad=True)
|
||||
w2 = Tensor.scaled_uniform(8, 4, requires_grad=True)
|
||||
|
||||
from tinygrad.nn import optim
|
||||
optimizer = optim.SGD([w1, w2], lr=0.01)
|
||||
|
||||
x = Tensor.rand(4, 16, device="AMD")
|
||||
h = (x @ w1).relu()
|
||||
y = h @ w2
|
||||
loss = y.sum()
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
self.assertIsNotNone(w1.grad)
|
||||
self.assertIsNotNone(w2.grad)
|
||||
|
||||
optimizer.step()
|
||||
# If we get here without GPU fault, the test passes
|
||||
|
||||
def test_conv_backward(self):
|
||||
"""Test conv2d backward pass"""
|
||||
Tensor.training = True
|
||||
x = Tensor.rand(2, 1, 8, 8, device="AMD", requires_grad=True)
|
||||
w = Tensor.rand(4, 1, 3, 3, device="AMD", requires_grad=True)
|
||||
y = x.conv2d(w)
|
||||
loss = y.sum()
|
||||
loss.backward()
|
||||
self.assertIsNotNone(x.grad)
|
||||
self.assertIsNotNone(w.grad)
|
||||
|
||||
def test_matmul_backward(self):
|
||||
"""Test matmul backward pass"""
|
||||
Tensor.training = True
|
||||
a = Tensor.rand(4, 8, device="AMD", requires_grad=True)
|
||||
b = Tensor.rand(8, 3, device="AMD", requires_grad=True)
|
||||
c = a @ b
|
||||
loss = c.sum()
|
||||
loss.backward()
|
||||
self.assertIsNotNone(a.grad)
|
||||
self.assertIsNotNone(b.grad)
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNAWMMA(unittest.TestCase):
|
||||
"""WMMA tensor core tests"""
|
||||
|
||||
def test_wmma_16x16_correctness(self):
|
||||
"""Test 16x16 WMMA correctness"""
|
||||
# Use uniform random values in [-0.5, 0.5) to avoid float16 overflow
|
||||
rng = np.random.default_rng(42)
|
||||
a = Tensor((rng.random((16, 16), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
b = Tensor((rng.random((16, 16), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
c_np = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
|
||||
np.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
|
||||
|
||||
def test_wmma_32x32_correctness(self):
|
||||
"""Test 32x32 WMMA correctness"""
|
||||
# Use uniform random values in [-0.5, 0.5) to avoid float16 overflow
|
||||
rng = np.random.default_rng(42)
|
||||
a = Tensor((rng.random((32, 32), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
b = Tensor((rng.random((32, 32), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
c_np = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
|
||||
np.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNAKernelGeneration(unittest.TestCase):
|
||||
"""Tests for kernel code generation"""
|
||||
|
||||
def test_wmma_kernel_runs(self):
|
||||
"""Test that WMMA kernel runs without error"""
|
||||
# If tensor cores are used, this matmul should run successfully
|
||||
a = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
b = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
c.realize()
|
||||
# Check result is reasonable
|
||||
self.assertEqual(c.shape, (16, 16))
|
||||
self.assertFalse(np.any(np.isnan(c.numpy())))
|
||||
|
||||
def test_larger_wmma_kernel_runs(self):
|
||||
"""Test that larger WMMA kernel (multiple tiles) runs"""
|
||||
a = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
b = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
c.realize()
|
||||
self.assertEqual(c.shape, (32, 32))
|
||||
self.assertFalse(np.any(np.isnan(c.numpy())))
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestRDNAVGPRLimits(unittest.TestCase):
|
||||
"""Tests documenting VGPR limits"""
|
||||
|
||||
@unittest.skip("64x64 WMMA exceeds VGPR limit (261 vs 256) - requires smaller tile size or improved register allocation")
|
||||
def test_wmma_64x64_vgpr_limit(self):
|
||||
"""64x64 WMMA - tests if VGPR limit is respected
|
||||
Currently skipped: 64x64 tiles require 261 VGPRs but RDNA3 limit is 256.
|
||||
Use N=16 for working WMMA tests.
|
||||
"""
|
||||
a = Tensor(np.random.randn(64, 64).astype(np.float16))
|
||||
b = Tensor(np.random.randn(64, 64).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
c_np = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
|
||||
c.realize()
|
||||
np.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
|
||||
|
||||
class TestWMMAVGPRUsage(unittest.TestCase):
|
||||
"""Tests for WMMA VGPR usage analysis - run without device"""
|
||||
|
||||
@unittest.skip("64x64 WMMA exceeds VGPR limit (261 vs 256) - known limitation")
|
||||
def test_64x64_wmma_vgpr_count(self):
|
||||
"""Test that 64x64 WMMA fits within 256 VGPR limit.
|
||||
|
||||
KNOWN LIMITATION: 64x64 WMMA tiles require 261 VGPRs but RDNA3 limit is 256.
|
||||
|
||||
Root cause analysis of VGPR usage:
|
||||
- 128 VGPRs for accumulators (16 output tiles x 8 floats each)
|
||||
- ~60 VGPRs for byte offset computation (SHL ops)
|
||||
- ~40 VGPRs for WMMA inputs (A and B matrices packed as half16)
|
||||
- Plus temps for index computation
|
||||
|
||||
Current optimizations applied:
|
||||
- CAST reuse: float32→half conversions reuse accumulator VGPRs
|
||||
- Deferred INDEX allocation: store addresses computed just-in-time
|
||||
- REG INDEX/LOAD: accumulator arrays don't allocate extra VGPRs
|
||||
|
||||
Future fix needed: defer SHL ops for byte offset computation to store time.
|
||||
"""
|
||||
import re
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
# Need TC=1 to enable tensor cores, TC_OPT=2 for padding support
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(64, 64).astype(np.float16))
|
||||
b = Tensor(np.random.randn(64, 64).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
|
||||
# Find the matmul kernel
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
has_reduce = any(u.op == Ops.REDUCE for u in item.ast.toposort())
|
||||
if has_reduce:
|
||||
# Use RDNARenderer for both full_rewrite and render
|
||||
# (rdna_matcher needs to be applied during full_rewrite to handle pointer CASTs)
|
||||
renderer = RDNARenderer('gfx1100')
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
|
||||
# Verify we have WMMA ops
|
||||
has_wmma = any(u.op == Ops.WMMA for u in uops)
|
||||
self.assertTrue(has_wmma, "Should have WMMA ops with USE_TC=1")
|
||||
|
||||
# Count WMMA and stores
|
||||
wmma_count = sum(1 for u in uops if u.op == Ops.WMMA)
|
||||
self.assertEqual(wmma_count, 16, f"Should have 16 WMMA ops for 64x64, got {wmma_count}")
|
||||
|
||||
# Render and check VGPR count
|
||||
asm = renderer.render(uops)
|
||||
|
||||
match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(match, "Should have .amdhsa_next_free_vgpr in metadata")
|
||||
vgpr_count = int(match.group(1))
|
||||
|
||||
# The VGPR limit is 256. This test documents the current state
|
||||
# and will pass once the register allocator is fixed.
|
||||
self.assertLessEqual(vgpr_count, 256,
|
||||
f"64x64 WMMA needs {vgpr_count} VGPRs but limit is 256. "
|
||||
f"Main issue: 128 store addresses computed upfront. "
|
||||
f"Fix: compute addresses just-in-time or reuse address VGPRs.")
|
||||
return
|
||||
|
||||
self.fail("Should find a kernel with REDUCE op")
|
||||
|
||||
def test_32x32_wmma_fits_in_vgpr_limit(self):
|
||||
"""Verify 32x32 WMMA fits within VGPR limit"""
|
||||
import re
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
b = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
has_reduce = any(u.op == Ops.REDUCE for u in item.ast.toposort())
|
||||
if has_reduce:
|
||||
# Use RDNARenderer for both full_rewrite and render
|
||||
renderer = RDNARenderer('gfx1100')
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
|
||||
has_wmma = any(u.op == Ops.WMMA for u in uops)
|
||||
self.assertTrue(has_wmma, "Should have WMMA ops")
|
||||
|
||||
asm = renderer.render(uops)
|
||||
|
||||
match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(match)
|
||||
vgpr_count = int(match.group(1))
|
||||
|
||||
# 32x32 should use fewer VGPRs (4 tiles instead of 16)
|
||||
# With perfect allocation: 4*8=32 accumulators + ~32 addresses + ~40 inputs = ~104
|
||||
self.assertLessEqual(vgpr_count, 256,
|
||||
f"32x32 WMMA needs {vgpr_count} VGPRs, should fit in 256")
|
||||
return
|
||||
|
||||
self.fail("Should find a kernel with REDUCE op")
|
||||
|
||||
class TestLookAheadPacking(unittest.TestCase):
|
||||
"""Tests for the look-ahead packing optimization - these run without device"""
|
||||
|
||||
def test_half16_const_packing(self):
|
||||
"""Test that half16 VECTORIZE with constants generates pack instructions"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Build a simple test: 16 scalar half constants -> half16 VECTORIZE
|
||||
# This tests the basic VECTORIZE packing logic (not look-ahead, but verifies packing works)
|
||||
half_vals = [UOp.const(dtypes.half, float(i)) for i in range(16)]
|
||||
vec = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(half_vals))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec,))
|
||||
|
||||
# Render and check for v_pack_b32_f16 instructions
|
||||
uops = list(sink.toposort())
|
||||
asm = renderer.render(uops)
|
||||
|
||||
# Should have v_pack_b32_f16 instructions for packing pairs of halfs
|
||||
pack_count = asm.count('v_pack_b32_f16')
|
||||
self.assertGreater(pack_count, 0, "Should have v_pack_b32_f16 instructions for half16 packing")
|
||||
# Should have 8 pack instructions (16 halfs / 2 per pack)
|
||||
self.assertEqual(pack_count, 8, f"Should have exactly 8 pack instructions, got {pack_count}")
|
||||
|
||||
def test_half16_load_packing_with_index(self):
|
||||
"""Test that half16 VECTORIZE with LOADs generates pack instructions and uses look-ahead"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Create a minimal kernel with half LOADs feeding half16 VECTORIZE
|
||||
# This simulates what WMMA needs for input data
|
||||
|
||||
# Create buffer argument using .ptr() method
|
||||
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0)
|
||||
|
||||
# Create index (workitem ID)
|
||||
ridx = UOp.special(32, "ridx0", dtype=dtypes.int)
|
||||
|
||||
# Create 16 LOADs at different offsets
|
||||
loads = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i)
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf.dtype, (buf, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads.append(load)
|
||||
|
||||
# Create half16 VECTORIZE
|
||||
vec = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec,))
|
||||
|
||||
uops = list(sink.toposort())
|
||||
asm = renderer.render(uops)
|
||||
|
||||
# Should have v_pack_b32_f16 instructions
|
||||
pack_count = asm.count('v_pack_b32_f16')
|
||||
self.assertGreater(pack_count, 0, "Should have v_pack_b32_f16 instructions")
|
||||
|
||||
# Should have 8 pack instructions total (16 halfs / 2 per pack)
|
||||
self.assertEqual(pack_count, 8, f"Should have exactly 8 pack instructions, got {pack_count}")
|
||||
|
||||
# Check that LOADs are generated (global_load_u16 for half precision)
|
||||
load_count = asm.count('global_load_u16')
|
||||
self.assertGreater(load_count, 0, "Should have global_load_u16 instructions")
|
||||
# Should have 16 loads (one per half element)
|
||||
self.assertEqual(load_count, 16, f"Should have 16 load instructions, got {load_count}")
|
||||
|
||||
def test_look_ahead_packing_pre_scan(self):
|
||||
"""Test that the pre-scan correctly identifies half16 VECTORIZE sources"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
# Create a minimal kernel with half LOADs feeding half16 VECTORIZE
|
||||
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0)
|
||||
ridx = UOp.special(32, "ridx0", dtype=dtypes.int)
|
||||
|
||||
# Create 16 LOADs
|
||||
loads = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i)
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf.dtype, (buf, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads.append(load)
|
||||
|
||||
vec = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec,))
|
||||
|
||||
uops = list(sink.toposort())
|
||||
|
||||
# Simulate the pre-scan logic from the renderer
|
||||
half16_vectorize_sources = {}
|
||||
for u in uops:
|
||||
if u.op is Ops.VECTORIZE and u.dtype.scalar() == dtypes.half and u.dtype.count == 16:
|
||||
for pos, src in enumerate(u.src):
|
||||
half16_vectorize_sources[src] = (u, pos)
|
||||
|
||||
# All 16 LOADs should be identified as half16 sources
|
||||
load_count = sum(1 for u in uops if u.op is Ops.LOAD and u in half16_vectorize_sources)
|
||||
self.assertEqual(load_count, 16, f"All 16 LOADs should be identified as half16 sources, got {load_count}")
|
||||
|
||||
def test_look_ahead_packing_is_interleaved(self):
|
||||
"""Test that pack instructions are interleaved with loads (not all at the end)"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.dtype import dtypes
|
||||
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Create a minimal kernel with half LOADs feeding half16 VECTORIZE
|
||||
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0)
|
||||
ridx = UOp.special(32, "ridx0", dtype=dtypes.int)
|
||||
|
||||
# Create 16 LOADs
|
||||
loads = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i)
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf.dtype, (buf, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads.append(load)
|
||||
|
||||
vec = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec,))
|
||||
|
||||
uops = list(sink.toposort())
|
||||
asm = renderer.render(uops)
|
||||
lines = asm.split('\n')
|
||||
|
||||
# Find line numbers of loads and packs
|
||||
load_lines = [i for i, l in enumerate(lines) if 'global_load_u16' in l]
|
||||
pack_lines = [i for i, l in enumerate(lines) if 'v_pack_b32_f16' in l]
|
||||
|
||||
self.assertGreater(len(load_lines), 0, "Should have load instructions")
|
||||
self.assertGreater(len(pack_lines), 0, "Should have pack instructions")
|
||||
|
||||
# The first pack should occur BEFORE the last load (interleaving)
|
||||
first_pack = min(pack_lines)
|
||||
last_load = max(load_lines)
|
||||
self.assertLess(first_pack, last_load,
|
||||
f"First pack (line {first_pack}) should be before last load (line {last_load}) for interleaving")
|
||||
|
||||
def test_vgpr_reuse_in_look_ahead_packing(self):
|
||||
"""Test that temp VGPRs are reused after packing (reducing register pressure)"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.dtype import dtypes
|
||||
import re
|
||||
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Create a minimal kernel with half LOADs feeding half16 VECTORIZE
|
||||
buf = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0)
|
||||
ridx = UOp.special(32, "ridx0", dtype=dtypes.int)
|
||||
|
||||
# Create 16 LOADs
|
||||
loads = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i)
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf.dtype, (buf, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads.append(load)
|
||||
|
||||
vec = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec,))
|
||||
|
||||
uops = list(sink.toposort())
|
||||
asm = renderer.render(uops)
|
||||
|
||||
# Extract VGPR count from metadata
|
||||
vgpr_match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(vgpr_match, "Should have .amdhsa_next_free_vgpr in metadata")
|
||||
vgpr_count = int(vgpr_match.group(1))
|
||||
|
||||
# With look-ahead packing and VGPR reuse:
|
||||
# - 8 VGPRs for the half16 destination range
|
||||
# - A few temp VGPRs for loading (reused)
|
||||
# - A few VGPRs for address computation
|
||||
# - 32 scratch VGPRs are allocated for potential 64-bit ops
|
||||
# Without reuse, we'd need 8 + 16 + 32 = 56 VGPRs minimum
|
||||
# With reuse, we should need fewer than that
|
||||
self.assertLess(vgpr_count, 56, f"VGPR count should be < 56 with reuse, got {vgpr_count}")
|
||||
|
||||
def test_multiple_half16_vectorizes(self):
|
||||
"""Test look-ahead packing with multiple half16 VECTORIZEs (like WMMA A and B inputs)"""
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.dtype import dtypes
|
||||
import re
|
||||
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Create two buffers (like A and B matrices)
|
||||
buf_a = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0)
|
||||
buf_b = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=1)
|
||||
ridx = UOp.special(32, "ridx0", dtype=dtypes.int)
|
||||
|
||||
# Create 16 LOADs from buf_a
|
||||
loads_a = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i)
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf_a.dtype, (buf_a, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads_a.append(load)
|
||||
|
||||
# Create 16 LOADs from buf_b
|
||||
loads_b = []
|
||||
for i in range(16):
|
||||
offset = UOp.const(dtypes.int, i + 100) # Different offsets
|
||||
idx = UOp(Ops.ADD, dtypes.int, (ridx, offset))
|
||||
index = UOp(Ops.INDEX, buf_b.dtype, (buf_b, idx))
|
||||
load = UOp(Ops.LOAD, dtypes.half, (index,))
|
||||
loads_b.append(load)
|
||||
|
||||
# Create two half16 VECTORIZEs
|
||||
vec_a = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads_a))
|
||||
vec_b = UOp(Ops.VECTORIZE, dtypes.half.vec(16), tuple(loads_b))
|
||||
sink = UOp(Ops.SINK, dtypes.void, (vec_a, vec_b))
|
||||
|
||||
uops = list(sink.toposort())
|
||||
asm = renderer.render(uops)
|
||||
|
||||
# Should have 16 pack instructions (8 for each half16)
|
||||
pack_count = asm.count('v_pack_b32_f16')
|
||||
self.assertEqual(pack_count, 16, f"Should have 16 pack instructions, got {pack_count}")
|
||||
|
||||
# Should have 32 loads total
|
||||
load_count = asm.count('global_load_u16')
|
||||
self.assertEqual(load_count, 32, f"Should have 32 load instructions, got {load_count}")
|
||||
|
||||
# Extract VGPR count
|
||||
vgpr_match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(vgpr_match, "Should have .amdhsa_next_free_vgpr in metadata")
|
||||
vgpr_count = int(vgpr_match.group(1))
|
||||
|
||||
# With look-ahead packing and VGPR reuse:
|
||||
# - 16 VGPRs for the two half16 destination ranges (8 each)
|
||||
# - A few temp VGPRs for loading (reused)
|
||||
# - 32 scratch VGPRs are allocated for potential 64-bit ops
|
||||
# Without reuse, we'd need 16 + 32 + 32 = 80 VGPRs minimum
|
||||
# With reuse, we should need fewer than that
|
||||
self.assertLess(vgpr_count, 80, f"VGPR count should be < 80 with reuse, got {vgpr_count}")
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestLookAheadPackingOnDevice(unittest.TestCase):
|
||||
"""Tests for look-ahead packing that require actual device execution"""
|
||||
|
||||
@unittest.skip("WMMA generation depends on scheduler decisions for matrix size")
|
||||
def test_wmma_generates_pack_instructions(self):
|
||||
"""Test that WMMA kernel generates v_pack_b32_f16 instructions"""
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.uop.ops import Ops
|
||||
|
||||
# Create matmul that uses WMMA
|
||||
a = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
b = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Find the WMMA kernel
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
uops = list(item.ast.toposort())
|
||||
has_wmma = any(u.op == Ops.WMMA for u in uops)
|
||||
if has_wmma:
|
||||
asm = renderer.render(uops)
|
||||
pack_count = asm.count('v_pack_b32_f16')
|
||||
self.assertGreater(pack_count, 0, "WMMA kernel should have v_pack_b32_f16 instructions")
|
||||
return
|
||||
|
||||
self.fail("Should find WMMA kernel in schedule")
|
||||
|
||||
@unittest.skip("Test needs adjustment for full_rewrite spec requirements")
|
||||
def test_look_ahead_packing_reduces_temp_regs(self):
|
||||
"""Test that look-ahead packing packs halfs immediately after load"""
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
# Create matmul that uses WMMA - need TC=1 and TC_OPT=2
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
b = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
renderer = RDNARenderer("gfx1100")
|
||||
|
||||
# Find and render the WMMA kernel
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
# Use full_rewrite to get lowered uops with WMMA
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
has_wmma = any(u.op == Ops.WMMA for u in uops)
|
||||
if has_wmma:
|
||||
asm = renderer.render(uops)
|
||||
lines = asm.split('\n')
|
||||
|
||||
# Check that pack instructions appear after loads (look-ahead packing)
|
||||
# The pattern should be: global_load_d16_b16 ... then v_pack_b32_f16
|
||||
load_lines = [i for i, l in enumerate(lines) if 'global_load_d16_b16' in l]
|
||||
pack_lines = [i for i, l in enumerate(lines) if 'v_pack_b32_f16' in l]
|
||||
|
||||
# Some pack instructions should appear interleaved with loads
|
||||
# (i.e., packing happens as soon as pairs are loaded, not all at the end)
|
||||
if load_lines and pack_lines:
|
||||
# Find first pack after a load
|
||||
first_pack = min(pack_lines)
|
||||
last_load = max(load_lines)
|
||||
# Pack should happen before all loads are done (interleaved)
|
||||
self.assertLess(first_pack, last_load,
|
||||
"Look-ahead packing should interleave pack instructions with loads")
|
||||
return
|
||||
|
||||
self.fail("Should find WMMA kernel in schedule")
|
||||
|
||||
@unittest.skipUnless(AMD_RDNA, "AMD_RDNA=1 required")
|
||||
class TestVGPRRegressions(unittest.TestCase):
|
||||
"""Regression tests for VGPR allocation optimizations"""
|
||||
|
||||
def test_16x16_wmma_on_device(self):
|
||||
"""Verify 16x16 WMMA matmul works correctly on device - regression test"""
|
||||
rng = np.random.default_rng(42)
|
||||
a = Tensor((rng.random((16, 16), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
b = Tensor((rng.random((16, 16), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
c_np = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
|
||||
c.realize()
|
||||
np.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
|
||||
|
||||
def test_32x32_wmma_on_device(self):
|
||||
"""Verify 32x32 WMMA matmul works correctly on device - regression test"""
|
||||
rng = np.random.default_rng(42)
|
||||
a = Tensor((rng.random((32, 32), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
b = Tensor((rng.random((32, 32), dtype=np.float32) - 0.5).astype(np.float16)).realize()
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
c_np = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
|
||||
c.realize()
|
||||
np.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
|
||||
|
||||
def test_deferred_store_index_detection(self):
|
||||
"""Verify that store-only INDEX ops are detected for deferred allocation"""
|
||||
import re
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
b = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
renderer = RDNARenderer('gfx1100')
|
||||
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
has_reduce = any(u.op == Ops.REDUCE for u in item.ast.toposort())
|
||||
if has_reduce:
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
|
||||
# Count STORE operations - there should be some
|
||||
store_count = sum(1 for u in uops if u.op == Ops.STORE)
|
||||
self.assertGreater(store_count, 0, "Should have STORE operations")
|
||||
|
||||
# Render and verify no register overflow
|
||||
asm = renderer.render(uops)
|
||||
match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(match)
|
||||
vgpr_count = int(match.group(1))
|
||||
self.assertLessEqual(vgpr_count, 256, f"Should fit in 256 VGPRs, got {vgpr_count}")
|
||||
return
|
||||
|
||||
self.fail("Should find kernel with REDUCE op")
|
||||
|
||||
def test_16x16_wmma_vgpr_count(self):
|
||||
"""Verify 16x16 WMMA uses reasonable VGPR count"""
|
||||
import re
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
b = Tensor(np.random.randn(16, 16).astype(np.float16))
|
||||
c = a.matmul(b, dtype=dtypes.float)
|
||||
|
||||
sched = c.schedule()
|
||||
renderer = RDNARenderer('gfx1100')
|
||||
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
has_reduce = any(u.op == Ops.REDUCE for u in item.ast.toposort())
|
||||
if has_reduce:
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
has_wmma = any(u.op == Ops.WMMA for u in uops)
|
||||
if not has_wmma:
|
||||
continue
|
||||
|
||||
asm = renderer.render(uops)
|
||||
match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(match)
|
||||
vgpr_count = int(match.group(1))
|
||||
|
||||
# 16x16 WMMA should use far fewer VGPRs than the limit
|
||||
# 1 WMMA tile = 8 accumulators, plus inputs/temps
|
||||
self.assertLess(vgpr_count, 128,
|
||||
f"16x16 WMMA should use <128 VGPRs, got {vgpr_count}")
|
||||
return
|
||||
|
||||
# If no WMMA kernel found, that's fine - test passes
|
||||
pass
|
||||
|
||||
def test_cast_reuse_optimization(self):
|
||||
"""Verify that CAST operations reuse accumulator VGPRs"""
|
||||
import re
|
||||
from tinygrad.codegen import full_rewrite
|
||||
from tinygrad.uop.ops import Ops
|
||||
from tinygrad.renderer.rdna_new import RDNARenderer
|
||||
from tinygrad.helpers import Context
|
||||
|
||||
with Context(TC=1, TC_OPT=2):
|
||||
a = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
b = Tensor(np.random.randn(32, 32).astype(np.float16))
|
||||
# Output as half to trigger float32->half CAST
|
||||
c = a.matmul(b, dtype=dtypes.float).cast(dtypes.half)
|
||||
|
||||
sched = c.schedule()
|
||||
renderer = RDNARenderer('gfx1100')
|
||||
|
||||
for item in sched:
|
||||
if hasattr(item, 'ast') and item.ast is not None:
|
||||
has_reduce = any(u.op == Ops.REDUCE for u in item.ast.toposort())
|
||||
if has_reduce:
|
||||
uops = full_rewrite(item.ast, renderer)
|
||||
asm = renderer.render(uops)
|
||||
|
||||
# Verify the kernel still fits in VGPR limit
|
||||
match = re.search(r'\.amdhsa_next_free_vgpr (\d+)', asm)
|
||||
self.assertIsNotNone(match)
|
||||
vgpr_count = int(match.group(1))
|
||||
self.assertLessEqual(vgpr_count, 256,
|
||||
f"32x32 WMMA with CAST should fit in 256 VGPRs, got {vgpr_count}")
|
||||
return
|
||||
|
||||
# If no kernel found, that's fine
|
||||
pass
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
+226
-73
@@ -13,15 +13,16 @@ from extra.assembly.rdna3.autogen import (
|
||||
# VOP1
|
||||
v_mov_b32_e32, v_cvt_f32_i32_e32, v_cvt_i32_f32_e32, v_cvt_f32_u32_e32, v_cvt_u32_f32_e32,
|
||||
v_cvt_f16_f32_e32, v_cvt_f32_f16_e32, v_rcp_f32_e32, v_sqrt_f32_e32,
|
||||
v_exp_f32_e32, v_log_f32_e32, v_trunc_f32_e32,
|
||||
v_exp_f32_e32, v_log_f32_e32, v_trunc_f32_e32, v_sin_f32_e32,
|
||||
v_cvt_f64_f32_e32, v_cvt_f32_f64_e32, v_cvt_f64_i32_e32, v_cvt_f64_u32_e32,
|
||||
v_cvt_i32_f64_e32, v_cvt_u32_f64_e32,
|
||||
v_cvt_i32_f64_e32, v_cvt_u32_f64_e32, v_trunc_f64_e32, v_floor_f64_e32,
|
||||
# VOP2
|
||||
v_add_f32_e32, v_sub_f32_e32, v_mul_f32_e32, v_and_b32_e32, v_or_b32_e32, v_xor_b32_e32,
|
||||
v_add_nc_u32_e32, v_sub_nc_u32_e32, v_lshlrev_b32_e32, v_lshrrev_b32_e32, v_ashrrev_i32_e32,
|
||||
v_max_f32_e32, v_max_i32_e32, v_max_u32_e32,
|
||||
# VOP3
|
||||
v_fma_f32, v_fma_f64, v_mad_u64_u32, v_mad_i64_i32, v_lshlrev_b64, v_mul_lo_u32, v_mul_hi_u32, v_bfe_u32,
|
||||
v_fma_f32, v_fma_f64, v_mad_u64_u32, v_mad_i64_i32, v_lshlrev_b64, v_lshrrev_b64, v_ashrrev_i64,
|
||||
v_mul_lo_u32, v_mul_hi_u32, v_bfe_u32, v_bfe_i32,
|
||||
v_add_co_u32, v_add_co_ci_u32_e32, v_cndmask_b32_e64, v_add_f64, v_mul_f64, v_sub_co_u32, v_sub_co_ci_u32_e32,
|
||||
v_cmp_lt_f32_e32, v_cmp_eq_f32_e32, v_cmp_neq_f32_e32, v_cmp_gt_f32_e32,
|
||||
v_cmp_lt_i32_e32, v_cmp_eq_i32_e32, v_cmp_ne_i32_e32, v_cmp_gt_i32_e32,
|
||||
@@ -38,91 +39,229 @@ from extra.assembly.rdna3.autogen import (
|
||||
)
|
||||
|
||||
# Helper for VOP2: src0 can be constant/literal, vsrc1 must be VGPR - swap for commutative ops
|
||||
def _sw(c, a, b):
|
||||
ar, br = c.get_reg(a), c.get_reg(b)
|
||||
def _sw(ctx, a, b):
|
||||
ar, br = ctx.get_reg(a), ctx.get_reg(b)
|
||||
return (br, ar) if isinstance(br, (int, float)) and not isinstance(ar, (int, float)) else (ar, br)
|
||||
|
||||
# Module-level PatternMatcher for simple ALU and CAST operations
|
||||
render_ops = PatternMatcher([
|
||||
# CAST: float32 <-> int32/uint32
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.int32),), name="x"), lambda c,x,a: [v_cvt_f32_i32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.uint32),), name="x"), lambda c,x,a: [v_cvt_f32_u32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_i32_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_u32_f32_e32(c.dst, c.get_reg(a))]),
|
||||
# CAST: float32 <-> small unsigned
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", (dtypes.uint8, dtypes.uint16)),), name="x"), lambda c,x,a: [v_cvt_f32_u32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_u32_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_i32_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.int32),), name="x"), lambda ctx,x,a: [v_cvt_f32_i32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.uint32),), name="x"), lambda ctx,x,a: [v_cvt_f32_u32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.float32),), name="x"), lambda ctx,x,a: [v_cvt_i32_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.float32),), name="x"), lambda ctx,x,a: [v_cvt_u32_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: float32 <-> small int
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", (dtypes.uint8, dtypes.uint16)),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_u32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.int8),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 8), v_cvt_f32_i32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.int16),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 16), v_cvt_f32_i32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.float32),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_u32_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.float32),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_i32_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: float16 <-> float32
|
||||
(UPat(Ops.CAST, dtypes.float16, (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_f16_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.float16),), name="x"), lambda c,x,a: [v_cvt_f32_f16_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float16, (UPat.var("a", dtypes.float32),), name="x"), lambda ctx,x,a: [v_cvt_f16_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.float16),), name="x"), lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: float16 -> ints (via f32)
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.float16),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a)), v_cvt_i32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.float16),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a)), v_cvt_u32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.float16),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a)), v_cvt_i32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.float16),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a)), v_cvt_u32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
# CAST: ints -> float16 (via f32)
|
||||
(UPat(Ops.CAST, dtypes.float16, (UPat.var("a", dtypes.ints),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_i32_e32(ctx.dst, ctx.get_reg(a)), v_cvt_f16_f32_e32(ctx.dst, ctx.dst)]),
|
||||
# CAST: bfloat16 <-> float32 (shift)
|
||||
(UPat(Ops.CAST, dtypes.bfloat16, (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_lshrrev_b32_e32(c.dst, 16, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.bfloat16),), name="x"), lambda c,x,a: [v_lshlrev_b32_e32(c.dst, 16, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.bfloat16, (UPat.var("a", dtypes.float32),), name="x"), lambda ctx,x,a: [v_lshrrev_b32_e32(ctx.dst, 16, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.bfloat16),), name="x"), lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a))]),
|
||||
# CAST: bfloat16 -> ints (via f32)
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.bfloat16),), name="x"),
|
||||
lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a)), v_cvt_i32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.bfloat16),), name="x"),
|
||||
lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a)), v_cvt_u32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.bfloat16),), name="x"),
|
||||
lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a)), v_cvt_i32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.bfloat16),), name="x"),
|
||||
lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a)), v_cvt_u32_f32_e32(ctx.dst, ctx.dst)]),
|
||||
# CAST: ints -> bfloat16 (via f32)
|
||||
(UPat(Ops.CAST, dtypes.bfloat16, (UPat.var("a", dtypes.ints),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_i32_e32(ctx.dst, ctx.get_reg(a)), v_lshrrev_b32_e32(ctx.dst, 16, ctx.dst)]),
|
||||
# CAST: float64 <-> float32
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.float32),), name="x"), lambda c,x,a: [v_cvt_f64_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.float64),), name="x"), lambda c,x,a: [v_cvt_f32_f64_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.float32),), name="x"), lambda ctx,x,a: [v_cvt_f64_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.float64),), name="x"), lambda ctx,x,a: [v_cvt_f32_f64_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: float64 <-> int32/uint32
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.int32),), name="x"), lambda c,x,a: [v_cvt_f64_i32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.uint32),), name="x"), lambda c,x,a: [v_cvt_f64_u32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.float64),), name="x"), lambda c,x,a: [v_cvt_i32_f64_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.float64),), name="x"), lambda c,x,a: [v_cvt_u32_f64_e32(c.dst, c.get_reg(a))]),
|
||||
# CAST: float64 <-> small unsigned
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", (dtypes.uint8, dtypes.uint16)),), name="x"), lambda c,x,a: [v_cvt_f64_u32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.float64),), name="x"), lambda c,x,a: [v_cvt_u32_f64_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.float64),), name="x"), lambda c,x,a: [v_cvt_i32_f64_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.int32),), name="x"), lambda ctx,x,a: [v_cvt_f64_i32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.uint32),), name="x"), lambda ctx,x,a: [v_cvt_f64_u32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.int32, (UPat.var("a", dtypes.float64),), name="x"), lambda ctx,x,a: [v_cvt_i32_f64_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.uint32, (UPat.var("a", dtypes.float64),), name="x"), lambda ctx,x,a: [v_cvt_u32_f64_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: float64 <-> small int
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", (dtypes.uint8, dtypes.uint16)),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f64_u32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.int8),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(v[ctx.dst.idx], ctx.get_reg(a), 0, 8), v_cvt_f64_i32_e32(ctx.dst, v[ctx.dst.idx])]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.int16),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(v[ctx.dst.idx], ctx.get_reg(a), 0, 16), v_cvt_f64_i32_e32(ctx.dst, v[ctx.dst.idx])]),
|
||||
(UPat(Ops.CAST, (dtypes.uint8, dtypes.uint16), (UPat.var("a", dtypes.float64),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_u32_f64_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, (dtypes.int8, dtypes.int16), (UPat.var("a", dtypes.float64),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_i32_f64_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: int64 -> smaller types (just take low 32 bits)
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", (dtypes.int64, dtypes.uint64)),), name="x"), lambda c,x,a: [v_mov_b32_e32(c.dst, c.get_reg(a))]),
|
||||
# CAST: small int <-> int32/uint32, small int <-> small int (just mov)
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", dtypes.ints),), name="x"), lambda c,x,a: [v_mov_b32_e32(c.dst, c.get_reg(a))]),
|
||||
# CAST: bool -> int (just move)
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", dtypes.bool),), name="x"), lambda c,x,a: [v_mov_b32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", (dtypes.int64, dtypes.uint64)),), name="x"), lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: int64/uint64 -> float32 (via float64: low + high*2^32, 2^32=0x41F0000000000000)
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.int64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(6), [v_cvt_f64_u32_e32(v[s:s+2], ctx.get_reg(a)),
|
||||
v_cvt_f64_i32_e32(v[s+2:s+4], v[ctx.get_reg(a).idx+1]),
|
||||
v_mov_b32_e32(v[s+4], 0), v_mov_b32_e32(v[s+5], 0x41F00000),
|
||||
v_mul_f64(v[s+2:s+4], v[s+4:s+6], v[s+2:s+4]),
|
||||
v_add_f64(v[s:s+2], v[s:s+2], v[s+2:s+4]), v_cvt_f32_f64_e32(ctx.dst, v[s:s+2])])[1]),
|
||||
(UPat(Ops.CAST, dtypes.float32, (UPat.var("a", dtypes.uint64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(6), [v_cvt_f64_u32_e32(v[s:s+2], ctx.get_reg(a)),
|
||||
v_cvt_f64_u32_e32(v[s+2:s+4], v[ctx.get_reg(a).idx+1]),
|
||||
v_mov_b32_e32(v[s+4], 0), v_mov_b32_e32(v[s+5], 0x41F00000),
|
||||
v_mul_f64(v[s+2:s+4], v[s+4:s+6], v[s+2:s+4]),
|
||||
v_add_f64(v[s:s+2], v[s:s+2], v[s+2:s+4]), v_cvt_f32_f64_e32(ctx.dst, v[s:s+2])])[1]),
|
||||
# CAST: int64/uint64 -> float64 (low + high*2^32, 2^32=0x41F0000000000000)
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.int64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(4), [v_cvt_f64_u32_e32(ctx.dst, ctx.get_reg(a)),
|
||||
v_cvt_f64_i32_e32(v[s:s+2], v[ctx.get_reg(a).idx+1]),
|
||||
v_mov_b32_e32(v[s+2], 0), v_mov_b32_e32(v[s+3], 0x41F00000),
|
||||
v_mul_f64(v[s:s+2], v[s+2:s+4], v[s:s+2]), v_add_f64(ctx.dst, ctx.dst, v[s:s+2])])[1]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.uint64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(4), [v_cvt_f64_u32_e32(ctx.dst, ctx.get_reg(a)),
|
||||
v_cvt_f64_u32_e32(v[s:s+2], v[ctx.get_reg(a).idx+1]),
|
||||
v_mov_b32_e32(v[s+2], 0), v_mov_b32_e32(v[s+3], 0x41F00000),
|
||||
v_mul_f64(v[s:s+2], v[s+2:s+4], v[s:s+2]), v_add_f64(ctx.dst, ctx.dst, v[s:s+2])])[1]),
|
||||
# CAST: float64 -> int64 (trunc, extract high*2^32 and low parts)
|
||||
# s[0:2]=trunc, s[2:4]=high_float, s[4:6]=trunc_copy for final subtract
|
||||
(UPat(Ops.CAST, dtypes.int64, (UPat.var("a", dtypes.float64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(6), ar:=ctx.get_reg(a), [
|
||||
v_trunc_f64_e32(v[s:s+2], ar), # Truncate to integer
|
||||
v_mov_b32_e32(v[s+4], v[s]), v_mov_b32_e32(v[s+5], v[s+1]), # Save trunc copy
|
||||
v_mov_b32_e32(v[s+2], 0), v_mov_b32_e32(v[s+3], 0x3DF00000), # 2^-32
|
||||
v_mul_f64(v[s+2:s+4], v[s:s+2], v[s+2:s+4]), # high = trunc / 2^32
|
||||
v_floor_f64_e32(v[s+2:s+4], v[s+2:s+4]), # floor(high)
|
||||
v_cvt_i32_f64_e32(v[ctx.dst.idx+1], v[s+2:s+4]), # high 32 bits
|
||||
v_mov_b32_e32(v[s], 0), v_mov_b32_e32(v[s+1], 0x41F00000), # 2^32
|
||||
v_mul_f64(v[s+2:s+4], v[s+2:s+4], v[s:s+2]), # high * 2^32
|
||||
v_mul_f64(v[s+2:s+4], -1.0, v[s+2:s+4]), # negate high*2^32
|
||||
v_add_f64(v[s:s+2], v[s+4:s+6], v[s+2:s+4]), # trunc - high*2^32 = low
|
||||
v_cvt_u32_f64_e32(ctx.dst, v[s:s+2])])[2]), # low 32 bits
|
||||
(UPat(Ops.CAST, dtypes.uint64, (UPat.var("a", dtypes.float64),), name="x"),
|
||||
lambda ctx,x,a: (s:=ctx.ra.get_scratch_vgpr(6), ar:=ctx.get_reg(a), [
|
||||
v_trunc_f64_e32(v[s:s+2], ar), # Truncate to integer
|
||||
v_mov_b32_e32(v[s+4], v[s]), v_mov_b32_e32(v[s+5], v[s+1]), # Save trunc copy
|
||||
v_mov_b32_e32(v[s+2], 0), v_mov_b32_e32(v[s+3], 0x3DF00000), # 2^-32
|
||||
v_mul_f64(v[s+2:s+4], v[s:s+2], v[s+2:s+4]), # high = trunc / 2^32
|
||||
v_floor_f64_e32(v[s+2:s+4], v[s+2:s+4]), # floor(high)
|
||||
v_cvt_u32_f64_e32(v[ctx.dst.idx+1], v[s+2:s+4]), # high 32 bits
|
||||
v_mov_b32_e32(v[s], 0), v_mov_b32_e32(v[s+1], 0x41F00000), # 2^32
|
||||
v_mul_f64(v[s+2:s+4], v[s+2:s+4], v[s:s+2]), # high * 2^32
|
||||
v_mul_f64(v[s+2:s+4], -1.0, v[s+2:s+4]), # negate high*2^32
|
||||
v_add_f64(v[s:s+2], v[s+4:s+6], v[s+2:s+4]), # trunc - high*2^32 = low
|
||||
v_cvt_u32_f64_e32(ctx.dst, v[s:s+2])])[2]), # low 32 bits
|
||||
# CAST: int8 -> larger types (sign extend)
|
||||
(UPat(Ops.CAST, (dtypes.int16, dtypes.uint16, dtypes.int32, dtypes.uint32), (UPat.var("a", dtypes.int8),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 8)]),
|
||||
# CAST: int16 -> 32-bit types (sign extend)
|
||||
(UPat(Ops.CAST, (dtypes.int32, dtypes.uint32), (UPat.var("a", dtypes.int16),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 16)]),
|
||||
# CAST: small signed int -> int64 (sign extend to 32-bit, then sign extend to 64-bit)
|
||||
(UPat(Ops.CAST, dtypes.int64, (UPat.var("a", dtypes.int8),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 8), v_ashrrev_i32_e32(v[ctx.dst.idx+1], 31, ctx.dst)]),
|
||||
(UPat(Ops.CAST, dtypes.int64, (UPat.var("a", dtypes.int16),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 16), v_ashrrev_i32_e32(v[ctx.dst.idx+1], 31, ctx.dst)]),
|
||||
# CAST: small signed int -> uint64 (sign extend to 32-bit, then zero extend high bits for unsigned reinterpret)
|
||||
(UPat(Ops.CAST, dtypes.uint64, (UPat.var("a", dtypes.int8),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 8), v_ashrrev_i32_e32(v[ctx.dst.idx+1], 31, ctx.dst)]),
|
||||
(UPat(Ops.CAST, dtypes.uint64, (UPat.var("a", dtypes.int16),), name="x"),
|
||||
lambda ctx,x,a: [v_bfe_i32(ctx.dst, ctx.get_reg(a), 0, 16), v_ashrrev_i32_e32(v[ctx.dst.idx+1], 31, ctx.dst)]),
|
||||
# CAST: int32 -> int64 (sign extend: copy sign bit to high 32 bits)
|
||||
(UPat(Ops.CAST, (dtypes.int64, dtypes.uint64), (UPat.var("a", dtypes.int32),), name="x"),
|
||||
lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a)), v_ashrrev_i32_e32(v[ctx.dst.idx+1], 31, ctx.get_reg(a))]),
|
||||
# CAST: uint32 -> int64/uint64 (zero extend: set high 32 bits to 0)
|
||||
(UPat(Ops.CAST, (dtypes.int64, dtypes.uint64), (UPat.var("a", dtypes.uint32),), name="x"),
|
||||
lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a)), v_mov_b32_e32(v[ctx.dst.idx+1], 0)]),
|
||||
# CAST: small unsigned int -> int64/uint64 (zero extend)
|
||||
(UPat(Ops.CAST, (dtypes.int64, dtypes.uint64), (UPat.var("a", (dtypes.uint8, dtypes.uint16)),), name="x"),
|
||||
lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a)), v_mov_b32_e32(v[ctx.dst.idx+1], 0)]),
|
||||
# CAST: small int <-> int32/uint32, small int <-> small int (just mov for unsigned or same-size)
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", dtypes.ints),), name="x"), lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# CAST: bool <-> int/float
|
||||
(UPat(Ops.CAST, (dtypes.int64, dtypes.uint64), (UPat.var("a", dtypes.bool),), name="x"),
|
||||
lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a)), v_mov_b32_e32(v[ctx.dst.idx+1], 0)]),
|
||||
(UPat(Ops.CAST, dtypes.ints, (UPat.var("a", dtypes.bool),), name="x"), lambda ctx,x,a: [v_mov_b32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.CAST, dtypes.bool, (UPat.var("a", dtypes.ints),), name="x"),
|
||||
lambda ctx,x,a: [v_cmp_ne_i32_e32(0, ctx.get_reg(a)), v_cndmask_b32_e64(ctx.dst, 0, 1, VCC_LO)]),
|
||||
(UPat(Ops.CAST, dtypes.bool, (UPat.var("a", dtypes.float32),), name="x"),
|
||||
lambda ctx,x,a: [v_cmp_neq_f32_e32(0.0, ctx.get_reg(a)), v_cndmask_b32_e64(ctx.dst, 0, 1, VCC_LO)]),
|
||||
(UPat(Ops.CAST, dtypes.bool, (UPat.var("a", dtypes.float16),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f16_e32(ctx.dst, ctx.get_reg(a)), v_cmp_neq_f32_e32(0.0, ctx.dst), v_cndmask_b32_e64(ctx.dst, 0, 1, VCC_LO)]),
|
||||
(UPat(Ops.CAST, dtypes.bool, (UPat.var("a", dtypes.float64),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f32_f64_e32(ctx.dst, ctx.get_reg(a)), v_cmp_neq_f32_e32(0.0, ctx.dst), v_cndmask_b32_e64(ctx.dst, 0, 1, VCC_LO)]),
|
||||
(UPat(Ops.CAST, dtypes.bool, (UPat.var("a", dtypes.bfloat16),), name="x"),
|
||||
lambda ctx,x,a: [v_lshlrev_b32_e32(ctx.dst, 16, ctx.get_reg(a)), v_cmp_neq_f32_e32(0.0, ctx.dst), v_cndmask_b32_e64(ctx.dst, 0, 1, VCC_LO)]),
|
||||
(UPat(Ops.CAST, dtypes.float64, (UPat.var("a", dtypes.bool),), name="x"),
|
||||
lambda ctx,x,a: [v_cvt_f64_u32_e32(ctx.dst, ctx.get_reg(a))]), # bool -> float64
|
||||
(UPat(Ops.CAST, dtypes.floats, (UPat.var("a", dtypes.bool),), name="x"), lambda ctx,x,a: [v_cvt_f32_u32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# ADD: float64, floats, int64, default to i32
|
||||
(UPat(Ops.ADD, dtype=dtypes.float64, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_add_f64(c.dst, c.get_reg(a), c.get_reg(b))]),
|
||||
(UPat(Ops.ADD, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_add_f32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.ADD, dtype=dtypes.float64, src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda ctx,x,a,b: [v_add_f64(ctx.dst, ctx.get_reg(a), ctx.get_reg(b))]),
|
||||
(UPat(Ops.ADD, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_add_f32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.ADD, dtype=(dtypes.int64, dtypes.uint64), src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda c,x,a,b: [v_add_co_u32(v[c.dst.idx], VCC_LO, v[c.get_reg(a).idx], v[c.get_reg(b).idx]),
|
||||
v_add_co_ci_u32_e32(v[c.dst.idx+1], v[c.get_reg(a).idx+1], v[c.get_reg(b).idx+1])]),
|
||||
(UPat(Ops.ADD, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_add_nc_u32_e32(c.dst, *_sw(c,a,b))]),
|
||||
lambda ctx,x,a,b: [v_add_co_u32(v[ctx.dst.idx], VCC_LO, v[ctx.get_reg(a).idx], v[ctx.get_reg(b).idx]),
|
||||
v_add_co_ci_u32_e32(v[ctx.dst.idx+1], v[ctx.get_reg(a).idx+1], v[ctx.get_reg(b).idx+1])]),
|
||||
(UPat(Ops.ADD, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_add_nc_u32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
# SUB: float64, floats, int64, default to i32
|
||||
(UPat(Ops.SUB, dtype=dtypes.float64, src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda c,x,a,b: [v_mul_f64(c.dst, -1.0, c.get_reg(b)), v_add_f64(c.dst, c.get_reg(a), c.dst)]),
|
||||
lambda ctx,x,a,b: [v_mul_f64(ctx.dst, -1.0, ctx.get_reg(b)), v_add_f64(ctx.dst, ctx.get_reg(a), ctx.dst)]),
|
||||
(UPat(Ops.SUB, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda c,x,a,b: [v_sub_f32_e32(c.dst, c.get_reg(a), c.get_reg(b))]),
|
||||
lambda ctx,x,a,b: [v_sub_f32_e32(ctx.dst, ctx.get_reg(a), ctx.get_reg(b))]),
|
||||
(UPat(Ops.SUB, dtype=(dtypes.int64, dtypes.uint64), src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda c,x,a,b: [v_sub_co_u32(v[c.dst.idx], VCC_LO, v[c.get_reg(a).idx], v[c.get_reg(b).idx]),
|
||||
v_sub_co_ci_u32_e32(v[c.dst.idx+1], v[c.get_reg(a).idx+1], v[c.get_reg(b).idx+1])]),
|
||||
(UPat(Ops.SUB, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_sub_nc_u32_e32(c.dst, c.get_reg(a), c.get_reg(b))]),
|
||||
lambda ctx,x,a,b: [v_sub_co_u32(v[ctx.dst.idx], VCC_LO, v[ctx.get_reg(a).idx], v[ctx.get_reg(b).idx]),
|
||||
v_sub_co_ci_u32_e32(v[ctx.dst.idx+1], v[ctx.get_reg(a).idx+1], v[ctx.get_reg(b).idx+1])]),
|
||||
(UPat(Ops.SUB, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_sub_nc_u32_e32(ctx.dst, ctx.get_reg(a), ctx.get_reg(b))]),
|
||||
# MUL: floats, ints (int64 complex - handled in emit_alu)
|
||||
(UPat(Ops.MUL, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_mul_f32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.MUL, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_mul_f32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.MUL, dtype=(dtypes.int32, dtypes.uint32, dtypes.int16, dtypes.uint16, dtypes.int8, dtypes.uint8),
|
||||
src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_mul_lo_u32(c.dst, *_sw(c,a,b))]),
|
||||
src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_mul_lo_u32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
# Bitwise: int only
|
||||
(UPat(Ops.AND, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_and_b32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.OR, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_or_b32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.XOR, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_xor_b32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.AND, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_and_b32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.OR, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_or_b32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.XOR, dtype=dtypes.ints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_xor_b32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
# SHL: int64, default to i32
|
||||
(UPat(Ops.SHL, dtype=(dtypes.int64, dtypes.uint64), src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda c,x,a,b: [v_lshlrev_b64(c.dst, c.get_reg(b), c.get_reg(a))]),
|
||||
(UPat(Ops.SHL, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_lshlrev_b32_e32(c.dst, c.get_reg(b), c.get_reg(a))]),
|
||||
lambda ctx,x,a,b: [v_lshlrev_b64(ctx.dst, ctx.get_reg(b), ctx.get_reg(a))]),
|
||||
(UPat(Ops.SHL, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_lshlrev_b32_e32(ctx.dst, ctx.get_reg(b), ctx.get_reg(a))]),
|
||||
# SHR: int64 signed, uint64 unsigned, default to 32-bit
|
||||
(UPat(Ops.SHR, dtype=dtypes.int64, src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda ctx,x,a,b: [v_ashrrev_i64(ctx.dst, ctx.get_reg(b), ctx.get_reg(a))]),
|
||||
(UPat(Ops.SHR, dtype=dtypes.uint64, src=(UPat.var("a"), UPat.var("b")), name="x"),
|
||||
lambda ctx,x,a,b: [v_lshrrev_b64(ctx.dst, ctx.get_reg(b), ctx.get_reg(a))]),
|
||||
# MAX: floats, signed ints, unsigned ints
|
||||
(UPat(Ops.MAX, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_max_f32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.MAX, dtype=dtypes.sints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_max_i32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.MAX, dtype=dtypes.uints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda c,x,a,b: [v_max_u32_e32(c.dst, *_sw(c,a,b))]),
|
||||
(UPat(Ops.MAX, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_max_f32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.MAX, dtype=dtypes.sints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_max_i32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
(UPat(Ops.MAX, dtype=dtypes.uints, src=(UPat.var("a"), UPat.var("b")), name="x"), lambda ctx,x,a,b: [v_max_u32_e32(ctx.dst, *_sw(ctx,a,b))]),
|
||||
# MULACC (FMA): float64, floats
|
||||
(UPat(Ops.MULACC, dtype=dtypes.float64, src=(UPat.var("a"), UPat.var("b"), UPat.var("d")), name="x"),
|
||||
lambda c,x,a,b,d: [v_fma_f64(c.dst, c.get_reg(a), c.get_reg(b), c.get_reg(d))]),
|
||||
lambda ctx,x,a,b,d: [v_fma_f64(ctx.dst, ctx.get_reg(a), ctx.get_reg(b), ctx.get_reg(d))]),
|
||||
(UPat(Ops.MULACC, dtype=dtypes.floats, src=(UPat.var("a"), UPat.var("b"), UPat.var("d")), name="x"),
|
||||
lambda c,x,a,b,d: [v_fma_f32(c.dst, c.get_reg(a), c.get_reg(b), c.get_reg(d))]),
|
||||
lambda ctx,x,a,b,d: [v_fma_f32(ctx.dst, ctx.get_reg(a), ctx.get_reg(b), ctx.get_reg(d))]),
|
||||
# Transcendental: float only
|
||||
(UPat(Ops.RECIPROCAL, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_rcp_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.SQRT, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_sqrt_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.EXP2, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_exp_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.LOG2, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_log_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.TRUNC, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_trunc_f32_e32(c.dst, c.get_reg(a))]),
|
||||
(UPat(Ops.RECIPROCAL, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_rcp_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.SQRT, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_sqrt_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.EXP2, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_exp_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.LOG2, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_log_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
(UPat(Ops.TRUNC, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_trunc_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# SIN: input should already be normalized by 1/(2π) via rdna_uops.py
|
||||
(UPat(Ops.SIN, dtype=dtypes.float32, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_sin_f32_e32(ctx.dst, ctx.get_reg(a))]),
|
||||
# NEG: floats vs ints
|
||||
(UPat(Ops.NEG, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_mul_f32_e32(c.dst, -1.0, c.get_reg(a))]),
|
||||
(UPat(Ops.NEG, dtype=dtypes.ints, src=(UPat.var("a"),), name="x"), lambda c,x,a: [v_sub_nc_u32_e32(c.dst, 0, c.get_reg(a))]),
|
||||
(UPat(Ops.NEG, dtype=dtypes.floats, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_mul_f32_e32(ctx.dst, -1.0, ctx.get_reg(a))]),
|
||||
(UPat(Ops.NEG, dtype=dtypes.ints, src=(UPat.var("a"),), name="x"), lambda ctx,x,a: [v_sub_nc_u32_e32(ctx.dst, 0, ctx.get_reg(a))]),
|
||||
])
|
||||
|
||||
|
||||
@@ -219,6 +358,25 @@ class RDNARenderer(Renderer):
|
||||
val = u.arg
|
||||
if val is Invalid: return 0 # Invalid index - return safe address (will be masked anyway)
|
||||
if isinstance(val, bool): return 1 if val else 0 # Handle bool before int (bool is subclass of int)
|
||||
# For 64-bit types, always load into register pair (can't use inline constants for 64-bit ADD)
|
||||
if u.dtype in (dtypes.int64, dtypes.uint64, dtypes.long, dtypes.ulong):
|
||||
reg = ra.alloc_vgpr_range(u, 2)
|
||||
lo = int(val) & 0xFFFFFFFF
|
||||
hi = (int(val) >> 32) & 0xFFFFFFFF
|
||||
code.append(v_mov_b32_e32(v[reg.idx], lo))
|
||||
code.append(v_mov_b32_e32(v[reg.idx + 1], hi))
|
||||
r[u] = reg
|
||||
return reg
|
||||
# Float64 constants: load as two 32-bit parts (IEEE 754 double)
|
||||
if u.dtype == dtypes.float64:
|
||||
reg = ra.alloc_vgpr_range(u, 2)
|
||||
bits64 = struct.unpack("Q", struct.pack("d", float(val)))[0]
|
||||
lo = bits64 & 0xFFFFFFFF
|
||||
hi = (bits64 >> 32) & 0xFFFFFFFF
|
||||
code.append(v_mov_b32_e32(v[reg.idx], lo))
|
||||
code.append(v_mov_b32_e32(v[reg.idx + 1], hi))
|
||||
r[u] = reg
|
||||
return reg
|
||||
if isinstance(val, float):
|
||||
if val in (0.0, 0.5, 1.0, 2.0, 4.0, -0.5, -1.0, -2.0, -4.0): return val
|
||||
# Convert inf/nan to hex representation
|
||||
@@ -226,15 +384,6 @@ class RDNARenderer(Renderer):
|
||||
val = struct.unpack("I", struct.pack("f", val))[0]
|
||||
elif isinstance(val, int) and -16 <= val <= 64: return val
|
||||
# Load literal constant into register
|
||||
# For 64-bit types, need to load both low and high 32 bits
|
||||
if u.dtype in (dtypes.int64, dtypes.uint64, dtypes.long, dtypes.ulong):
|
||||
reg = ra.alloc_vgpr_range(u, 2)
|
||||
lo = val & 0xFFFFFFFF
|
||||
hi = (val >> 32) & 0xFFFFFFFF
|
||||
code.append(v_mov_b32_e32(v[reg.idx], lo))
|
||||
code.append(v_mov_b32_e32(v[reg.idx + 1], hi))
|
||||
r[u] = reg
|
||||
return reg
|
||||
reg = ra.alloc_vgpr(u)
|
||||
code.append(v_mov_b32_e32(reg, val))
|
||||
r[u] = reg
|
||||
@@ -301,13 +450,17 @@ class RDNARenderer(Renderer):
|
||||
code.append(v_lshrrev_b32_e32(dst, b - 32, high_reg) if is_unsigned else v_ashrrev_i32_e32(dst, b - 32, high_reg))
|
||||
else:
|
||||
code.append(v_lshrrev_b32_e32(dst, b, a) if is_unsigned else v_ashrrev_i32_e32(dst, b, a))
|
||||
elif op is Ops.SIN:
|
||||
raise NotImplementedError("SIN requires input normalization by 1/(2π) - needs to be lowered in rdna_uops.py")
|
||||
elif op in (Ops.CMPLT, Ops.CMPEQ, Ops.CMPNE):
|
||||
emit_cmp(op, u.src[0].dtype, dst, a, b)
|
||||
elif op is Ops.WHERE:
|
||||
cond, true_val, false_val = srcs[0], srcs[1], srcs[2]
|
||||
code.extend([v_cmp_ne_i32_e32(0, cond), v_cndmask_b32_e64(dst, false_val, true_val, VCC_LO)])
|
||||
code.append(v_cmp_ne_i32_e32(0, cond))
|
||||
if dtype == dtypes.float64:
|
||||
# For float64: select both low and high 32-bit parts
|
||||
code.append(v_cndmask_b32_e64(v[dst.idx], v[false_val.idx], v[true_val.idx], VCC_LO))
|
||||
code.append(v_cndmask_b32_e64(v[dst.idx+1], v[false_val.idx+1], v[true_val.idx+1], VCC_LO))
|
||||
else:
|
||||
code.append(v_cndmask_b32_e64(dst, false_val, true_val, VCC_LO))
|
||||
elif op is Ops.IDIV:
|
||||
# Integer division using floating-point approximation
|
||||
# quotient = trunc(float(a) * rcp(float(b)))
|
||||
|
||||
@@ -67,7 +67,7 @@ def lower_umod(a: UOp, b: UOp) -> UOp:
|
||||
|
||||
def lower_idiv(a: UOp, b: UOp) -> UOp:
|
||||
"""Lower signed 32-bit division using unsigned division on absolute values."""
|
||||
zero, one = UOp.const(dtypes.int32, 0), UOp.const(dtypes.int32, 1)
|
||||
zero = UOp.const(dtypes.int32, 0)
|
||||
a_neg, b_neg = a.alu(Ops.CMPLT, zero), b.alu(Ops.CMPLT, zero)
|
||||
a_abs = UOp(Ops.WHERE, dtypes.int32, (a_neg, zero - a, a)).bitcast(dtypes.uint32)
|
||||
b_abs = UOp(Ops.WHERE, dtypes.int32, (b_neg, zero - b, b)).bitcast(dtypes.uint32)
|
||||
@@ -183,6 +183,9 @@ def _lower_bf16_to_f16(x: UOp) -> UOp:
|
||||
"""bfloat16 -> float16: go through float32."""
|
||||
return x.src[0].cast(dtypes.float32).cast(dtypes.float16)
|
||||
|
||||
# *** Cast lowerings: multi-step casts go via intermediate types ***
|
||||
_small_ints = (dtypes.int8, dtypes.int16, dtypes.uint8, dtypes.uint16)
|
||||
|
||||
# Pattern matcher for RDNA3-specific rewrites
|
||||
# NOTE: By the time rdna_matcher runs, gated loads have already been created by devectorize
|
||||
# (WHERE+LOAD -> LOAD(INDEX(buf, idx, gate), alt)). We don't need to do that transformation here.
|
||||
@@ -196,8 +199,51 @@ rdna_matcher = PatternMatcher([
|
||||
(UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat(dtype=dtypes.half),), name="x"), _lower_f16_to_bf16),
|
||||
(UPat(Ops.CAST, dtype=dtypes.float16, src=(UPat(dtype=dtypes.bfloat16),), name="x"), _lower_bf16_to_f16),
|
||||
(UPat(Ops.CAST, dtype=dtypes.half, src=(UPat(dtype=dtypes.bfloat16),), name="x"), _lower_bf16_to_f16),
|
||||
# small ints <-> float16/bfloat16 via float32
|
||||
(UPat(Ops.CAST, dtype=_small_floats, src=(UPat(dtype=_small_ints),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat(Ops.CAST, dtype=_small_ints, src=(UPat(dtype=_small_floats),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
# int32/uint32 <-> float16/bfloat16 via float32
|
||||
(UPat(Ops.CAST, dtype=_small_floats, src=(UPat(dtype=(dtypes.int32, dtypes.uint32)),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat(Ops.CAST, dtype=(dtypes.int32, dtypes.uint32), src=(UPat(dtype=_small_floats),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
# int64/uint64 <-> float32 via float64
|
||||
(UPat(Ops.CAST, dtype=dtypes.float32, src=(UPat(dtype=(dtypes.int64, dtypes.uint64)),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float64).cast(dtypes.float32)),
|
||||
(UPat(Ops.CAST, dtype=(dtypes.int64, dtypes.uint64), src=(UPat(dtype=dtypes.float32),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float64).cast(x.dtype)),
|
||||
# int64/uint64 <-> float16/bfloat16 via float64 -> float32
|
||||
(UPat(Ops.CAST, dtype=_small_floats, src=(UPat(dtype=(dtypes.int64, dtypes.uint64)),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float64).cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat(Ops.CAST, dtype=(dtypes.int64, dtypes.uint64), src=(UPat(dtype=_small_floats),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.float32).cast(dtypes.float64).cast(x.dtype)),
|
||||
# small ints <-> float64 via float32
|
||||
(UPat(Ops.CAST, dtype=dtypes.float64, src=(UPat(dtype=_small_ints),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(dtypes.float64)),
|
||||
(UPat(Ops.CAST, dtype=_small_ints, src=(UPat(dtype=dtypes.float64),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
# float16/bfloat16 <-> float64 via float32
|
||||
(UPat(Ops.CAST, dtype=dtypes.float64, src=(UPat(dtype=_small_floats),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(dtypes.float64)),
|
||||
(UPat(Ops.CAST, dtype=_small_floats, src=(UPat(dtype=dtypes.float64),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
# bool <-> float16/bfloat16 via float32
|
||||
(UPat(Ops.CAST, dtype=_small_floats, src=(UPat(dtype=dtypes.bool),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(x.dtype)),
|
||||
(UPat(Ops.CAST, dtype=dtypes.bool, src=(UPat(dtype=_small_floats),), name="x"), lambda x: x.src[0].cast(dtypes.float32).cast(dtypes.bool)),
|
||||
# bool <-> int64/uint64 (need to handle 64-bit extension)
|
||||
(UPat(Ops.CAST, dtype=(dtypes.int64, dtypes.uint64), src=(UPat(dtype=dtypes.bool),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.int32).cast(x.dtype)),
|
||||
(UPat(Ops.CAST, dtype=dtypes.bool, src=(UPat(dtype=(dtypes.int64, dtypes.uint64)),), name="x"),
|
||||
lambda x: x.src[0].cast(dtypes.int32).cast(dtypes.bool)),
|
||||
# float64 comparisons: lower to float32 for now (VOP3 CMP needs special handling)
|
||||
(UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.float64), UPat.var("b")), name="x"),
|
||||
lambda x, a, b: UOp(Ops.CMPLT, dtypes.bool, (a.cast(dtypes.float32), b.cast(dtypes.float32)))),
|
||||
(UPat(Ops.CMPEQ, src=(UPat.var("a", dtypes.float64), UPat.var("b")), name="x"),
|
||||
lambda x, a, b: UOp(Ops.CMPEQ, dtypes.bool, (a.cast(dtypes.float32), b.cast(dtypes.float32)))),
|
||||
(UPat(Ops.CMPNE, src=(UPat.var("a", dtypes.float64), UPat.var("b")), name="x"),
|
||||
lambda x, a, b: UOp(Ops.CMPNE, dtypes.bool, (a.cast(dtypes.float32), b.cast(dtypes.float32)))),
|
||||
# devectorize ALU operations - RDNA doesn't have vector float ALU
|
||||
(UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name="alu"), no_vectorized_alu),
|
||||
# SIN: normalize input by 1/(2π) for v_sin_f32 (expects [0,1) -> [0,2π))
|
||||
(UPat(Ops.SIN, dtype=dtypes.float32, src=(UPat.var("x"),), name="u"),
|
||||
lambda u, x: None if u.tag == "normalized" else # skip already normalized
|
||||
UOp(Ops.SIN, dtypes.float32, (x * UOp.const(dtypes.float32, 0.15915494309189535),)).rtag("normalized")),
|
||||
# Fix fast_idiv output when shift >= 32 (needs 64-bit multiply)
|
||||
# Pattern: (x * const) >> shift for unsigned
|
||||
(UPat(Ops.SHR, src=(UPat(Ops.MUL, src=(UPat.var("x"), UPat.cvar("c"))), UPat.cvar("shift"))), _fix_fast_idiv_unsigned),
|
||||
|
||||
Reference in New Issue
Block a user