diff --git a/examples/vectoradd_py/submission_mojo.py b/examples/vectoradd_py/submission_mojo.py new file mode 100644 index 000000000..7bf2714f4 --- /dev/null +++ b/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 diff --git a/src/kernelbot/cogs/leaderboard_cog.py b/src/kernelbot/cogs/leaderboard_cog.py index 01d5c5838..8432f79b5 100644 --- a/src/kernelbot/cogs/leaderboard_cog.py +++ b/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" diff --git a/src/libkernelbot/task.py b/src/libkernelbot/task.py index ddd819451..7d5cb63a4 100644 --- a/src/libkernelbot/task.py +++ b/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: diff --git a/src/runners/modal_runner.py b/src/runners/modal_runner.py index a457a44ca..e0af5014d 100644 --- a/src/runners/modal_runner.py +++ b/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 \\n' | " "nvcc -std=c++17 -x cu -c - -o /tmp/cublasdx-smoke.o && " diff --git a/tests/test_modal.py b/tests/test_modal.py index 241edfb58..2c31fe557 100644 --- a/tests/test_modal.py +++ b/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(