From 4f6e3f679d3dcf64b7b7f7799aa31f6a92d33c41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Fri, 6 Nov 2026 15:09:54 +0300 Subject: [PATCH 1/3] Add Mojo implementation for vector addition in PMPP v2 --- .../solutions/correct/submission_mojo.py | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py diff --git a/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py new file mode 100644 index 000000000..6684d0488 --- /dev/null +++ b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py @@ -0,0 +1,110 @@ +#!POPCORN leaderboard vectoradd_v2 +import sys +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_v2() abi("C") -> PythonObject: + try: + var mb = PythonModuleBuilder("mojo_vecadd_v2") + _ = ( + 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_v2: {e}") +""" + +_pkg = Path("/tmp/mojo_vecadd_v2") +_pkg.mkdir(exist_ok=True) +(_pkg / "mojo_vecadd_v2.mojo").write_text(MOJO_SOURCE) +sys.path.insert(0, str(_pkg)) +import mojo_vecadd_v2 # noqa: E402 + +_LAUNCH = mojo_vecadd_v2.Runner().launch + + +def custom_kernel(data: input_t) -> output_t: + A, B, C = data + assert A.is_cuda and B.is_cuda and C.is_cuda, "Input/output tensors must be on GPU" + assert A.shape == B.shape, "Input tensors must have the same shape" + assert C.shape == A.shape, "Output tensor and input tensors must have the same shape" + assert A.dtype == torch.float16 and B.dtype == torch.float16 and C.dtype == torch.float16, "Input/output tensors must be float16" + _LAUNCH(A.data_ptr(), B.data_ptr(), C.data_ptr(), A.numel()) + return C From c47680c300f89ff26a69909aebca9d33c1aac514 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 7 Nov 2026 17:00:32 +0300 Subject: [PATCH 2/3] Load the v2 Mojo solution through `importlib` from a unique temporary directory with `tempfile`. --- .../vectoradd_py/solutions/correct/submission_mojo.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py index 6684d0488..79f670680 100644 --- a/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py +++ b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py @@ -1,5 +1,7 @@ #!POPCORN leaderboard vectoradd_v2 +import importlib import sys +import tempfile from pathlib import Path import mojo.importer # noqa: F401 @@ -91,11 +93,10 @@ def PyInit_mojo_vecadd_v2() abi("C") -> PythonObject: abort(t"Failed to initialize mojo_vecadd_v2: {e}") """ -_pkg = Path("/tmp/mojo_vecadd_v2") -_pkg.mkdir(exist_ok=True) +_pkg = Path(tempfile.mkdtemp(prefix="mojo_vecadd_v2_")) (_pkg / "mojo_vecadd_v2.mojo").write_text(MOJO_SOURCE) sys.path.insert(0, str(_pkg)) -import mojo_vecadd_v2 # noqa: E402 +mojo_vecadd_v2 = importlib.import_module("mojo_vecadd_v2") _LAUNCH = mojo_vecadd_v2.Runner().launch From 41ba1330c00bb9ba898c4b330687bacbd06ee038 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=2E=20Tolga=20Cang=C3=B6z?= Date: Sat, 7 Nov 2026 18:06:41 +0300 Subject: [PATCH 3/3] Refactor Mojo loading in `vectoradd_v2` to use the current file's directory instead of a temporary directory --- .../vectoradd_py/solutions/correct/submission_mojo.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py index 79f670680..e56fc6625 100644 --- a/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py +++ b/problems/pmpp_v2/vectoradd_py/solutions/correct/submission_mojo.py @@ -1,7 +1,5 @@ #!POPCORN leaderboard vectoradd_v2 import importlib -import sys -import tempfile from pathlib import Path import mojo.importer # noqa: F401 @@ -93,9 +91,7 @@ def PyInit_mojo_vecadd_v2() abi("C") -> PythonObject: abort(t"Failed to initialize mojo_vecadd_v2: {e}") """ -_pkg = Path(tempfile.mkdtemp(prefix="mojo_vecadd_v2_")) -(_pkg / "mojo_vecadd_v2.mojo").write_text(MOJO_SOURCE) -sys.path.insert(0, str(_pkg)) +(Path(__file__).resolve().parent / "mojo_vecadd_v2.mojo").write_text(MOJO_SOURCE) mojo_vecadd_v2 = importlib.import_module("mojo_vecadd_v2") _LAUNCH = mojo_vecadd_v2.Runner().launch