RoPE kernel for latency test: issue #166 (issue #94), standalone 8-core (4x2)

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<HBM>
+    } : 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<HBM>
+    } : 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<HBM>
+    } : 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<HBM>
+    } : 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
+  }
+}

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": [

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."""

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.
 

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"
+            )