diff --git a/extra/assembly/rdna3/asm.py b/extra/assembly/rdna3/asm.py index 65d57a63c7..26aa68b535 100644 --- a/extra/assembly/rdna3/asm.py +++ b/extra/assembly/rdna3/asm.py @@ -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') diff --git a/test/test_rdna_renderer.py b/test/test_rdna_renderer.py deleted file mode 100644 index 1d17b74c3e..0000000000 --- a/test/test_rdna_renderer.py +++ /dev/null @@ -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() diff --git a/tinygrad/renderer/rdna_new.py b/tinygrad/renderer/rdna_new.py index 35380d1058..79548d8193 100644 --- a/tinygrad/renderer/rdna_new.py +++ b/tinygrad/renderer/rdna_new.py @@ -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))) diff --git a/tinygrad/renderer/rdna_uops.py b/tinygrad/renderer/rdna_uops.py index ed95d4676a..1e9c41a115 100644 --- a/tinygrad/renderer/rdna_uops.py +++ b/tinygrad/renderer/rdna_uops.py @@ -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),