cecilia/feature/nki-mul2

benchmarks/operators/mul2/impl_nki.py

@@ -0,0 +1,61 @@
+import torch
+
+try:
+    import neuronxcc.nki as nki
+    import neuronxcc.nki.language as nl
+    import neuronxcc.nki.isa as nisa
+    PMAX = nl.tile_size.pmax
+except ImportError:
+    nki = None
+
+if nki is not None:
+    @nki.jit
+    def mul2_kernel(a_input):
+        num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX
+
+        free_tile_size = 16384
+
+        num_free_blocks = (a_input.shape[1] + free_tile_size - 1) // free_tile_size
+
+        hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm)
+
+        for i in range(num_blocks):
+            offset = i * PMAX
+
+            partition_index = nl.arange(PMAX)[:, None]
+
+            mask_p = partition_index < (a_input.shape[0] - offset)
+
+            for j in range(num_free_blocks):
+                free_offset = j * free_tile_size
+
+                free_dim_index = nl.arange(free_tile_size)[None, :]
+
+                mask_f = free_dim_index < (a_input.shape[1] - free_offset)
+
+                mask = mask_p & mask_f
+
+                a_tile = nl.load(a_input[offset + partition_index, free_offset + free_dim_index], mask=mask)
+
+                result_tile = nl.multiply(a_tile, 2, mask=mask)
+
+                nl.store(hbm_result_tile[offset + partition_index, free_offset + free_dim_index], value=result_tile, mask=mask)
+
+        return hbm_result_tile
+
+def run(x: torch.Tensor, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor:
+    n = x.numel()
+
+    free_dim = (n + (PMAX -1)) // PMAX
+    padded_size = PMAX * free_dim
+
+    if padded_size > n:
+        x = torch.nn.functional.pad(x, (0, padded_size - n))
+
+    x_2d = x.reshape(PMAX, free_dim)
+
+    result = mul2_kernel(x_2d)
+    return result.reshape(-1)[:n]
+
+def get_last_config() -> dict | None:
+    return None