diff --git a/examples/latency/rope_fwd_4x2.mlir b/examples/latency/rope_fwd_4x2.mlir new file mode 100644 index 0000000..7ef6524 --- /dev/null +++ b/examples/latency/rope_fwd_4x2.mlir @@ -0,0 +1,146 @@ +// RoPE Forward — Standalone Rotary Position Embedding +// +// Formulation (half-layout, LLaMA convention): +// y[:, 0:D/2] = x[:, 0:D/2] * cos - x[:, D/2:D] * sin +// y[:, D/2:D] = x[:, 0:D/2] * sin + x[:, D/2:D] * cos +// +// Reference model: LLaMA-3-8B / Granite-8B (H_q=32, H_kv=8, D=128, S=4096) +// Total head-lanes: 40 (32 Q + 8 K, both rotated) +// +// Grid (4, 2): dim0=seq partitions (4), dim1=head groups (2) +// - x/out: PARTITIONED along both dims (block-contiguous) +// - cos/sin: PARTITIONED along dim0 (seq), REPLICATED along dim1 (heads) +// - cos/sin depend only on seq position → shared across all heads +// - No inter-core communication (embarrassingly parallel) +// +// Per-core: 20 heads (= 40 total / 2 grid_dim1) × 1024 positions × 128 dim +// Loop structure: seq-tile loop (4 iter, TILE_SEQ=256) → head loop (20 iter) +// cos/sin [256, 64] loaded once per seq-tile, reused across all 20 heads +// +// Arithmetic intensity: ~0.77 FLOPs/byte (memory-bound, vector-unit workload) +// Peak working set per inner iteration: 192 KB (6 × [256,64] × 2B) + +module { + func.func @rope_fwd_kernel( + %x_ptr: index, // input [H=40, S=4096, D=128] flattened to [163840, 128] f16 + %cos_ptr: index, // precomputed cos [S=4096, D/2=64] f16 + %sin_ptr: index, // precomputed sin [S=4096, D/2=64] f16 + %out_ptr: index // output [H=40, S=4096, D=128] flattened to [163840, 128] f16 + ) attributes {grid = [4, 2]} { + + // --- Tile IDs and offsets --- + %pid_s, %pid_h = ktdp.get_compute_tile_id : index, index + + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c64 = arith.constant 64 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %c4096 = arith.constant 4096 : index + + // Loop bounds (compile-time constants) + %NUM_SEQ_TILES = arith.constant 4 : index // 1024 positions per core / 256 TILE_SEQ + %HEADS_PER_CORE = arith.constant 20 : index // (32 Q + 8 K) heads / 2 grid_dim1 + + %seq_offset = arith.muli %pid_s, %c1024 : index + %head_offset = arith.muli %pid_h, %HEADS_PER_CORE : index + + // --- Construct memory views --- + %x_view = ktdp.construct_memory_view %x_ptr, sizes: [163840, 128], strides: [128, 1] { + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 163839 >= 0, d1 >= 0, -d1 + 127 >= 0)>, + memory_space = #ktdp.spyre_memory_space + } : memref<163840x128xf16> + + %cos_view = ktdp.construct_memory_view %cos_ptr, sizes: [4096, 64], strides: [64, 1] { + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + memory_space = #ktdp.spyre_memory_space + } : memref<4096x64xf16> + + %sin_view = ktdp.construct_memory_view %sin_ptr, sizes: [4096, 64], strides: [64, 1] { + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 4095 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + memory_space = #ktdp.spyre_memory_space + } : memref<4096x64xf16> + + %out_view = ktdp.construct_memory_view %out_ptr, sizes: [163840, 128], strides: [128, 1] { + coordinate_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 163839 >= 0, d1 >= 0, -d1 + 127 >= 0)>, + memory_space = #ktdp.spyre_memory_space + } : memref<163840x128xf16> + + // --- OUTER LOOP: seq-tile loop (TILE_SEQ=256) --- + scf.for %t = %c0 to %NUM_SEQ_TILES step %c1 { + + %tile_offset = arith.muli %t, %c256 : index + %cos_row = arith.addi %seq_offset, %tile_offset : index + + // Load cos/sin [256, 64] once per seq-tile, reused across all 20 heads + %cos_acc = ktdp.construct_access_tile %cos_view[%cos_row, %c0] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<4096x64xf16> -> !ktdp.access_tile<256x64xindex> + + %cos_tile = ktdp.load %cos_acc : !ktdp.access_tile<256x64xindex> -> tensor<256x64xf16> + + %sin_acc = ktdp.construct_access_tile %sin_view[%cos_row, %c0] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<4096x64xf16> -> !ktdp.access_tile<256x64xindex> + + %sin_tile = ktdp.load %sin_acc : !ktdp.access_tile<256x64xindex> -> tensor<256x64xf16> + + // --- INNER LOOP: head loop --- + scf.for %h = %c0 to %HEADS_PER_CORE step %c1 { + + // Row in flattened [H*S, D]: (head_offset + h) * S + seq_offset + tile_offset + %h_abs = arith.addi %head_offset, %h : index + %row_base = arith.muli %h_abs, %c4096 : index + %row = arith.addi %row_base, %cos_row : index + + // Load x_first [256, 64] — first half of head_dim + %x_first_acc = ktdp.construct_access_tile %x_view[%row, %c0] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<163840x128xf16> -> !ktdp.access_tile<256x64xindex> + + %x_first = ktdp.load %x_first_acc : !ktdp.access_tile<256x64xindex> -> tensor<256x64xf16> + + // Load x_second [256, 64] — second half of head_dim + %x_second_acc = ktdp.construct_access_tile %x_view[%row, %c64] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<163840x128xf16> -> !ktdp.access_tile<256x64xindex> + + %x_second = ktdp.load %x_second_acc : !ktdp.access_tile<256x64xindex> -> tensor<256x64xf16> + + // Compute y_first = x_first * cos - x_second * sin + %tmp1 = arith.mulf %x_first, %cos_tile : tensor<256x64xf16> + %tmp2 = arith.mulf %x_second, %sin_tile : tensor<256x64xf16> + %y_first = arith.subf %tmp1, %tmp2 : tensor<256x64xf16> + + // Store y_first [256, 64] + %y_first_acc = ktdp.construct_access_tile %out_view[%row, %c0] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<163840x128xf16> -> !ktdp.access_tile<256x64xindex> + + ktdp.store %y_first, %y_first_acc : tensor<256x64xf16>, !ktdp.access_tile<256x64xindex> + + // Compute y_second = x_first * sin + x_second * cos + %tmp3 = arith.mulf %x_first, %sin_tile : tensor<256x64xf16> + %tmp4 = arith.mulf %x_second, %cos_tile : tensor<256x64xf16> + %y_second = arith.addf %tmp3, %tmp4 : tensor<256x64xf16> + + // Store y_second [256, 64] + %y_second_acc = ktdp.construct_access_tile %out_view[%row, %c64] { + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 + 255 >= 0, d1 >= 0, -d1 + 63 >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + } : memref<163840x128xf16> -> !ktdp.access_tile<256x64xindex> + + ktdp.store %y_second, %y_second_acc : tensor<256x64xf16>, !ktdp.access_tile<256x64xindex> + + scf.yield + } + scf.yield + } + return + } +} diff --git a/tests/conftest.py b/tests/conftest.py index 8a530e2..e3e29d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -397,6 +397,17 @@ }, ], # --------------------------------------------------------------------------- + # RoPE forward (Issue #166) + # --------------------------------------------------------------------------- + "rope_fwd_kernel": [ + { + "path": "latency/rope_fwd_4x2.mlir", + # LLaMA-8B / Granite-8B: H=40, S=4096, D=128, grid=[4,2] + # All dimensions baked into MLIR; no scalar kwargs needed. + "execute_kwargs": {}, + }, + ], + # --------------------------------------------------------------------------- # FFN-SwiGLU example (Issue #77) # --------------------------------------------------------------------------- "ffn_swiglu": [ diff --git a/tests/mlir_frontend/test_examples_adapt.py b/tests/mlir_frontend/test_examples_adapt.py index f5abf2e..8bb93d9 100644 --- a/tests/mlir_frontend/test_examples_adapt.py +++ b/tests/mlir_frontend/test_examples_adapt.py @@ -20,6 +20,7 @@ TestSdpaExecution as _TestSdpaExecution, TestPagedAttentionExecution as _TestPagedAttentionExecution, TestScalarBroadcastExecution as _TestScalarBroadcastExecution, + TestRoPEExecution as _TestRoPEExecution, ) @@ -68,3 +69,7 @@ class TestPagedAttentionAdapt(MLIRFrontendInterpMixin, _TestPagedAttentionExecut class TestScalarBroadcastAdapt(MLIRFrontendInterpMixin, _TestScalarBroadcastExecution): """Scalar broadcast (rank-0 collapse) via MLIRFrontendParser.""" + + +class TestRoPEAdapt(MLIRFrontendInterpMixin, _TestRoPEExecution): + """RoPE tests via MLIRFrontendParser.""" diff --git a/tests/test_examples.py b/tests/test_examples.py index f872d31..c0676f5 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1077,6 +1077,116 @@ def _prepare_and_seed(grid_shape): np.testing.assert_allclose(result, expected, rtol=1e-3, atol=1e-3) +class TestRoPEExecution(InterpreterTestMixin): + """End-to-end execution of rope_fwd_4x2.mlir. + + Tests the standalone RoPE kernel (half-layout, LLaMA convention): + y[:, 0:D/2] = x[:, 0:D/2] * cos - x[:, D/2:D] * sin + y[:, D/2:D] = x[:, 0:D/2] * sin + x[:, D/2:D] * cos + + LLaMA-8B / Granite-8B: H_total=40, S=4096, D=128, grid=[4,2] + """ + + H_TOTAL = 40 + S = 4096 + D = 128 + D_HALF = 64 + + @staticmethod + def _make_cos_sin_tables(S, D_half): + """Generate cos/sin tables with standard RoPE frequencies.""" + freqs = 10000.0 ** (-np.arange(D_half).astype(np.float64) * 2.0 / (D_half * 2)) + positions = np.arange(S, dtype=np.float64) + angles = np.outer(positions, freqs) + cos_table = np.cos(angles).astype(np.float16) + sin_table = np.sin(angles).astype(np.float16) + return cos_table, sin_table + + @staticmethod + def _rope_reference(x_flat, cos_table, sin_table, H, S, D): + """NumPy reference for half-layout RoPE.""" + D_half = D // 2 + x_3d = x_flat.reshape(H, S, D).astype(np.float32) + cos_f32 = cos_table.astype(np.float32)[np.newaxis, :, :] # [1, S, D/2] + sin_f32 = sin_table.astype(np.float32)[np.newaxis, :, :] + + x_first = x_3d[:, :, :D_half] + x_second = x_3d[:, :, D_half:] + + y_first = x_first * cos_f32 - x_second * sin_f32 + y_second = x_first * sin_f32 + x_second * cos_f32 + + return np.concatenate([y_first, y_second], axis=-1).reshape(H * S, D).astype(np.float16) + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("rope_fwd_kernel")) + def test_rope_fwd_correctness(self, path, func_name, entry): + """Run RoPE on 8 cores, verify output matches NumPy reference.""" + interp = self._make_interp() + interp.load(path) + + x_ptr, cos_ptr, sin_ptr, out_ptr = interp.arg_names(func_name) + sizes = interp.tensor_input_output_sizes(func_name) + + rows, cols = sizes[x_ptr]["shape"] + assert rows == self.H_TOTAL * self.S + assert cols == self.D + + cos_rows, cos_cols = sizes[cos_ptr]["shape"] + assert cos_rows == self.S + assert cos_cols == self.D_HALF + + rng = np.random.default_rng(42) + x = rng.standard_normal((rows, cols)).astype(np.float16) + cos_table, sin_table = self._make_cos_sin_tables(self.S, self.D_HALF) + out = np.zeros((rows, cols), dtype=np.float16) + + outputs = interp.execute_function(func_name, **{ + x_ptr: x, + cos_ptr: cos_table, + sin_ptr: sin_table, + out_ptr: out, + }) + + result = outputs[out_ptr] + expected = self._rope_reference(x, cos_table, sin_table, self.H_TOTAL, self.S, self.D) + + assert result.shape == expected.shape + assert not np.any(np.isnan(result)), "output contains NaN" + assert not np.any(np.isinf(result)), "output contains Inf" + + np.testing.assert_allclose( + result, + expected, + rtol=1e-2, + atol=1e-2, + err_msg="RoPE output does not match NumPy reference", + ) + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("rope_fwd_kernel")) + def test_rope_fwd_zero_input(self, path, func_name, entry): + """Zero input should produce zero output (rotation of origin is origin).""" + interp = self._make_interp() + interp.load(path) + + x_ptr, cos_ptr, sin_ptr, out_ptr = interp.arg_names(func_name) + sizes = interp.tensor_input_output_sizes(func_name) + rows, cols = sizes[x_ptr]["shape"] + + x = np.zeros((rows, cols), dtype=np.float16) + cos_table, sin_table = self._make_cos_sin_tables(self.S, self.D_HALF) + out = np.zeros((rows, cols), dtype=np.float16) + + outputs = interp.execute_function(func_name, **{ + x_ptr: x, + cos_ptr: cos_table, + sin_ptr: sin_table, + out_ptr: out, + }) + + result = outputs[out_ptr] + np.testing.assert_allclose(result, 0.0, atol=1e-5) + + class TestNestedYieldExecution(InterpreterTestMixin): """End-to-end execution of nested_yield.ktir. diff --git a/tests/test_latency.py b/tests/test_latency.py index 48e279b..3c13b2c 100644 --- a/tests/test_latency.py +++ b/tests/test_latency.py @@ -2868,3 +2868,97 @@ def test_memory_scales_with_bandwidth(self, path, func_name, entry, hbm_bw): expected_ratio = 1.0 / hbm_bw actual_ratio = scaled_mem / baseline_mem assert actual_ratio == pytest.approx(expected_ratio, rel=1e-3) + + +# --------------------------------------------------------------------------- +# RoPE forward latency — memory-dominated, vector-unit workload +# --------------------------------------------------------------------------- + + +def _run_rope(path, func_name, entry, cfg, trace=False): + """Run RoPE kernel on 8 cores and return report.""" + interp = KTIRInterpreter(latency_config=cfg, trace_latency=trace) + interp.load(path) + + sizes = interp.tensor_input_output_sizes(func_name) + x_ptr, cos_ptr, sin_ptr, out_ptr = interp.arg_names(func_name) + + rows, cols = sizes[x_ptr]["shape"] + cos_rows, cos_cols = sizes[cos_ptr]["shape"] + + rng = np.random.default_rng(42) + x = rng.standard_normal((rows, cols)).astype(np.float16) + + freqs = 10000.0 ** (-np.arange(cos_cols, dtype=np.float64) * 2.0 / (cos_cols * 2)) + positions = np.arange(cos_rows, dtype=np.float64) + angles = np.outer(positions, freqs) + cos_table = np.cos(angles).astype(np.float16) + sin_table = np.sin(angles).astype(np.float16) + out = np.zeros((rows, cols), dtype=np.float16) + + interp.execute_function(func_name, **{ + x_ptr: x, + cos_ptr: cos_table, + sin_ptr: sin_table, + out_ptr: out, + }) + return interp.get_latency_report() + + +class TestRoPELatency: + """Latency tests for standalone RoPE kernel (issue #166). + + Three properties: + 1. Memory-bound: data movement dominates, scales with HBM bandwidth + 2. SIMD compute: vector-unit workload, scales with SIMD width + 3. No communication: embarrassingly parallel across all 8 cores + """ + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("rope_fwd_kernel")) + @pytest.mark.parametrize("hbm_bw", [0.256, 0.512]) + def test_memory_bound(self, path, func_name, entry, hbm_bw): + """Data movement: memory-dominated bottleneck, AI below ridge, scales with BW.""" + baseline_cfg = HardwareConfig(hbm_bandwidth_tb_s=0.128) + scaled_cfg = HardwareConfig(hbm_bandwidth_tb_s=hbm_bw) + + baseline = _run_rope(path, func_name, entry, baseline_cfg) + scaled = _run_rope(path, func_name, entry, scaled_cfg) + + # Memory-bound classification (at realistic BW, memory dominates) + assert baseline.bottleneck == "memory" + assert baseline.counters[0].memory_cycles > baseline.counters[0].compute_cycles + rf = baseline.roofline() + assert rf["core_AI"] < rf["ridge"] + + # Memory cycles scale as 1/bandwidth + expected_ratio = 0.128 / hbm_bw + actual_ratio = scaled.counters[0].memory_cycles / baseline.counters[0].memory_cycles + assert actual_ratio == pytest.approx(expected_ratio, rel=1e-3) + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("rope_fwd_kernel")) + @pytest.mark.parametrize("simd", [32, 64]) + def test_simd_compute(self, path, func_name, entry, simd): + """Computation: pure vector-unit (simd), no matmul, scales with SIMD width.""" + baseline_cfg = HardwareConfig(simd_elements_per_cycle=64) + scaled_cfg = HardwareConfig(simd_elements_per_cycle=simd) + + baseline = _run_rope(path, func_name, entry, baseline_cfg) + scaled = _run_rope(path, func_name, entry, scaled_cfg) + + # Dominant compute unit is simd (mulf/subf/addf, no linalg.matmul) + rf = baseline.roofline() + assert rf["core_dominant_unit"] == "simd" + + # Compute cycles scale as 1/simd_width + expected_ratio = 64.0 / simd + actual_ratio = scaled.counters[0].compute_cycles / baseline.counters[0].compute_cycles + assert actual_ratio == pytest.approx(expected_ratio, rel=1e-2) + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("rope_fwd_kernel")) + def test_no_communication(self, path, func_name, entry): + """Communication: embarrassingly parallel — zero inter-core traffic.""" + report = _run_rope(path, func_name, entry, HardwareConfig()) + for cid, counters in report.counters.items(): + assert counters.comm_cycles == 0, ( + f"Core {cid} has {counters.comm_cycles} comm cycles, expected 0" + )