Propose to add Mojo🔥-based kernels

examples/vectoradd_py/submission_mojo.py

@@ -0,0 +1,104 @@
+#!POPCORN leaderboard vectoradd_py-dev
+import importlib
+from pathlib import Path
+
+import mojo.importer  # noqa: F401
+import torch
+from task import input_t, output_t
+
+MOJO_SOURCE = """
+from std.math import ceildiv
+from std.memory import Pointer
+from std.os import abort
+from std.python import PythonObject
+from std.python.bindings import PythonModuleBuilder
+from max.gpu import global_idx
+from max.gpu.host import DeviceContext
+from layout import RowMajorLayout, TileTensor, row_major
+
+comptime BLOCK = 256
+comptime Layout1D = RowMajorLayout[Int]
+
+
+def add_kernel(
+    a: TileTensor[.float16, Layout1D, ImmUntrackedOrigin],
+    b: TileTensor[.float16, Layout1D, ImmUntrackedOrigin],
+    c: TileTensor[.float16, Layout1D, MutUntrackedOrigin],
+    n: Int32,
+):
+    var idx = Int32(global_idx.x)
+    if idx < n:
+        c[idx] = a[idx] + b[idx]
+
+
+@fieldwise_init
+struct Runner(Movable, Writable):
+    var ctx: DeviceContext
+    var func: type_of(DeviceContext().compile_function[add_kernel]())
+
+    def write_repr_to(self, mut writer: Some[Writer]):
+        writer.write("Runner()")
+
+    @staticmethod
+    def py_init(out self: Self, args: PythonObject, kwargs: PythonObject) raises:
+        var ctx = DeviceContext()
+        var func = ctx.compile_function[add_kernel]()
+        self = Self(ctx^, func^)
+
+    @staticmethod
+    def launch(
+        self_ptr: Pointer[Self, MutAnyOrigin],
+        a_py: PythonObject,
+        b_py: PythonObject,
+        c_py: PythonObject,
+        n_py: PythonObject,
+    ) raises:
+        var n = Int(py=n_py)
+        var a = TileTensor(
+            Pointer[Float16, ImmUntrackedOrigin](unsafe_from_address=Int(py=a_py)),
+            row_major(n),
+        )
+        var b = TileTensor(
+            Pointer[Float16, ImmUntrackedOrigin](unsafe_from_address=Int(py=b_py)),
+            row_major(n),
+        )
+        var c = TileTensor(
+            Pointer[Float16, MutUntrackedOrigin](unsafe_from_address=Int(py=c_py)),
+            row_major(n),
+        )
+        self_ptr[].ctx.enqueue_function(
+            self_ptr[].func,
+            a,
+            b,
+            c,
+            Int32(n),
+            grid_dim=ceildiv(n, BLOCK),
+            block_dim=BLOCK,
+        )
+
+
+@export
+def PyInit_mojo_vecadd_example() abi("C") -> PythonObject:
+    try:
+        var mb = PythonModuleBuilder("mojo_vecadd_example")
+        _ = (
+            mb.add_type[Runner]("Runner")
+            .def_py_init[Runner.py_init]()
+            .def_method[Runner.launch]("launch")
+        )
+        return mb.finalize()
+    except e:
+        abort(t"Failed to initialize mojo_vecadd_example: {e}")
+"""
+
+(Path(__file__).resolve().parent / "mojo_vecadd_example.mojo").write_text(MOJO_SOURCE)
+mojo_vecadd_example = importlib.import_module("mojo_vecadd_example")
+
+_LAUNCH = mojo_vecadd_example.Runner().launch
+
+
+def custom_kernel(data: input_t) -> output_t:
+    A, B = data
+    C = torch.empty_like(A)
+    _LAUNCH(A.data_ptr(), B.data_ptr(), C.data_ptr(), A.numel())
+    return C

src/kernelbot/cogs/leaderboard_cog.py

@@ -239,7 +239,14 @@ async def lang_autocomplete(
 
 
 def add_header_to_template(lang: str, code: str, lb: LeaderboardItem):
-    comment_char = {"CUDA": "//", "Python": "#", "Triton": "#", "HIP": "#", "CuteDSL": "#"}[lang]
+    comment_char = {
+        "CUDA": "//",
+        "Python": "#",
+        "Triton": "#",
+        "HIP": "#",
+        "CuteDSL": "#",
+        "Mojo": "#",
+    }[lang]
 
     description_comment = [f"{comment_char} > {line}" for line in lb["description"].splitlines()]
     header = f"""
@@ -548,7 +555,14 @@ async def get_task_template(
                 return
 
             template = add_header_to_template(lang, templates[lang], leaderboard_item)
-            ext = {"CUDA": "cu", "Python": "py", "Triton": "py", "HIP": "py", "CuteDSL": "py"}
+            ext = {
+                "CUDA": "cu",
+                "Python": "py",
+                "Triton": "py",
+                "HIP": "py",
+                "CuteDSL": "py",
+                "Mojo": "py",
+            }
             file_name = f"{leaderboard_name}.{ext[lang]}"
             file = discord.File(fp=StringIO(template), filename=file_name)
             message = f"**Starter code for leaderboard `{leaderboard_name}`**\n"

src/libkernelbot/task.py

@@ -163,7 +163,7 @@ def make_task_definition(yaml_file: str | Path) -> LeaderboardDefinition:  # noq
     # load template files
     templates = {}
     for lang, source in raw.get("templates", {}).items():
-        assert lang in ["CUDA", "Python", "Triton", "HIP", "CuteDSL"]
+        assert lang in ["CUDA", "Python", "Triton", "HIP", "CuteDSL", "Mojo"]
         templates[lang] = (root / source).read_text()
 
     if templates:

src/runners/modal_runner.py

@@ -68,6 +68,7 @@
     .uv_pip_install(
         "tinygrad~=0.10",
         "helion",
+        "max",
     )
     # nvidia cuda packages
     .uv_pip_install(
@@ -108,8 +109,9 @@
     })
     .run_commands(
         "python -m pip check",
-        'python -c "import cuda.tile, nvmath"',
+        'python -c "import cuda.tile, nvmath, mojo.importer, max"',
         "tileiras --version",
+        "mojo --version",
         (
             "printf '#include <cublasdx.hpp>\\n' | "
             "nvcc -std=c++17 -x cu -c - -o /tmp/cublasdx-smoke.o && "

tests/test_modal.py

@@ -197,6 +197,7 @@ def modal_deployment(project_root: Path):
     [
         ("vectoradd_py", "submission_cuda_inline.py"),
         ("vectoradd_py", "submission_triton.py"),
+        ("vectoradd_py", "submission_mojo.py"),
     ],
 )
 async def test_modal_launcher_python_script(