Writing Float64 into numpy via UnsafePointer corrupts when also reading another float64 buffer (Mojo 1.0b3)

Curious · August 5, 2026, 2:26am

We build a Python extension with PythonModuleBuilder and need a fast path that

fills preallocated numpy float64/int64 buffers from Mojo (COO weight build).

Pattern that fails on `Mojo 1.0.0b3.dev2026080406` + numpy 2.5.1:


# Read fx via UnsafePointer, write row/col/weight via UnsafePointer

var x = fx_p[unsafe_offset=i]

row_p[unsafe_offset=i] = Int64(i)

col_p[unsafe_offset=i] = Int64(i * 10)

w_p[unsafe_offset=i] = x * 0.5

At n=4096, row/col match expected values, but weight[:4] is garbage (for example, bit patterns that look like pointers, or zeros). Both serial and parallelize versions fail.

A single-buffer fill such as:

arr[i] = Float64(i)

works correctly.

Full repro attached:

numpy_ptr_repro.mojo

"""Minimal repro: Mojo UnsafePointer ↔ numpy buffer issues.

Context (aether remapper): production stencil build stays on NumPy because
writing COO through ``UnsafePointer`` into preallocated numpy arrays was
observed to SEGV or corrupt (weights ↔ col indices / garbage) under Mojo
1.0b3, including serial loops. Workaround: Python list append + ``np.asarray``.

This file has several cases. On Mojo 1.0.0b3.dev2026080406, simple serial
fills pass; run all cases and report which fail on your build.

Forum ask: supported way to write Mojo Float64/Int64 into an existing
numpy ndarray from a ``PythonModuleBuilder`` extension (serial + parallelize)?

Run::

    cd scripts/mojo_numpy_pointer_repro
    rm -rf __mojocache__
    pixi run python run_repro.py
"""

from std.os import abort
from std.python import Python, PythonObject
from std.python.bindings import PythonModuleBuilder
from std.memory import UnsafePointer
from max.algorithm import parallelize


def _f64_ptr(arr: PythonObject) raises -> UnsafePointer[Float64, MutUntrackedOrigin]:
    return UnsafePointer[Float64, MutUntrackedOrigin](
        unsafe_from_address=Int(py=arr.ctypes.data)
    )


def _i64_ptr(arr: PythonObject) raises -> UnsafePointer[Int64, MutUntrackedOrigin]:
    return UnsafePointer[Int64, MutUntrackedOrigin](
        unsafe_from_address=Int(py=arr.ctypes.data)
    )


def fill_f64(n_obj: PythonObject) raises -> PythonObject:
    """Serial: arr[i] = Float64(i)."""
    var np = Python.import_module("numpy")
    var n = Int(py=n_obj)
    var arr = np.empty(n, dtype=np.float64)
    var p = _f64_ptr(arr)
    var i = 0
    while i < n:
        p[unsafe_offset=i] = Float64(i)
        i += 1
    return arr


def fill_i64(n_obj: PythonObject) raises -> PythonObject:
    """Serial: arr[i] = Int64(i)."""
    var np = Python.import_module("numpy")
    var n = Int(py=n_obj)
    var arr = np.empty(n, dtype=np.int64)
    var p = _i64_ptr(arr)
    var i = 0
    while i < n:
        p[unsafe_offset=i] = Int64(i)
        i += 1
    return arr


def fill_f64_parallel(n_obj: PythonObject) raises -> PythonObject:
    """parallelize: arr[i] = Float64(i)."""
    var np = Python.import_module("numpy")
    var os = Python.import_module("os")
    var n = Int(py=n_obj)
    var arr = np.empty(n, dtype=np.float64)
    var p = _f64_ptr(arr)
    var workers = Int(py=os.cpu_count())
    if workers < 1:
        workers = 1
    if workers > n:
        workers = n
    if workers > 64:
        workers = 64

    @parameter
    def work(i: Int):
        p[unsafe_offset=i] = Float64(i)

    parallelize[work](n, workers)
    return arr


def stencil_like(n_obj: PythonObject) raises -> PythonObject:
    """Closer to aether COO build: read fx, write row/col/weight via pointers.

    Expected for each i in [0, n):
      row[i] = i
      col[i] = i * 10
      weight[i] = fx[i] * 0.5   (fx[i] was set to Float64(i) in Python)
    """
    var np = Python.import_module("numpy")
    var n = Int(py=n_obj)
    var fx = np.arange(n, dtype=np.float64)
    var row = np.empty(n, dtype=np.int64)
    var col = np.empty(n, dtype=np.int64)
    var weight = np.empty(n, dtype=np.float64)

    var fx_p = _f64_ptr(fx)
    var row_p = _i64_ptr(row)
    var col_p = _i64_ptr(col)
    var w_p = _f64_ptr(weight)

    var i = 0
    while i < n:
        var x = fx_p[unsafe_offset=i]
        row_p[unsafe_offset=i] = Int64(i)
        col_p[unsafe_offset=i] = Int64(i * 10)
        w_p[unsafe_offset=i] = x * 0.5
        i += 1

    var out = Python.dict()
    out["row"] = row
    out["col"] = col
    out["weight"] = weight
    return out


def stencil_like_parallel(n_obj: PythonObject) raises -> PythonObject:
    """Same as stencil_like but with parallelize (aether's old hot path)."""
    var np = Python.import_module("numpy")
    var os = Python.import_module("os")
    var n = Int(py=n_obj)
    var fx = np.arange(n, dtype=np.float64)
    var row = np.empty(n, dtype=np.int64)
    var col = np.empty(n, dtype=np.int64)
    var weight = np.empty(n, dtype=np.float64)

    var fx_p = _f64_ptr(fx)
    var row_p = _i64_ptr(row)
    var col_p = _i64_ptr(col)
    var w_p = _f64_ptr(weight)

    var workers = Int(py=os.cpu_count())
    if workers < 1:
        workers = 1
    if workers > n:
        workers = n
    if workers > 64:
        workers = 64

    @parameter
    def work(i: Int):
        var x = fx_p[unsafe_offset=i]
        row_p[unsafe_offset=i] = Int64(i)
        col_p[unsafe_offset=i] = Int64(i * 10)
        w_p[unsafe_offset=i] = x * 0.5

    parallelize[work](n, workers)

    var out = Python.dict()
    out["row"] = row
    out["col"] = col
    out["weight"] = weight
    return out


def fill_f64_via_python_index(n_obj: PythonObject) raises -> PythonObject:
    """Control: PythonObject indexing (slow, usually correct)."""
    var np = Python.import_module("numpy")
    var n = Int(py=n_obj)
    var arr = np.empty(n, dtype=np.float64)
    var i = 0
    while i < n:
        arr[i] = Float64(i)
        i += 1
    return arr


@export
def PyInit_numpy_ptr_repro() abi("C") -> PythonObject:
    try:
        var m = PythonModuleBuilder("numpy_ptr_repro")
        m.def_function[fill_f64]("fill_f64")
        m.def_function[fill_i64]("fill_i64")
        m.def_function[fill_f64_parallel]("fill_f64_parallel")
        m.def_function[stencil_like]("stencil_like")
        m.def_function[stencil_like_parallel]("stencil_like_parallel")
        m.def_function[fill_f64_via_python_index]("fill_f64_via_python_index")
        return m.finalize()
    except e:
        abort(String("failed to create numpy_ptr_repro: ", e))

run_repro.py

#!/usr/bin/env python3
"""Driver for numpy_ptr_repro.mojo — print expected vs got for forum posts."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import mojo.importer  # noqa: E402, F401
import numpy_ptr_repro as ker  # noqa: E402


def check(name: str, got: np.ndarray, expected: np.ndarray) -> bool:
    ok = np.array_equal(got, expected) and got.dtype == expected.dtype
    print(f"\n=== {name} (n={got.size}, dtype={got.dtype}) ===")
    print("expected[:8]:", expected[:8])
    print("got[:8]:     ", got[:8])
    if not ok:
        mism = np.where(got != expected)[0]
        print(f"MISMATCH: {mism.size}/{got.size} indices differ; first={mism[:10]}")
        if got.size and mism.size:
            i = int(mism[0])
            print(f"  got[{i}]={got[i]!r} expected[{i}]={expected[i]!r}")
            print(f"  got.view(np.uint64)[{i}]=0x{np.asarray(got).view(np.uint64)[i]:016x}")
    else:
        print("OK")
    return ok


def check_stencil(name: str, n: int, fn) -> bool:
    print(f"\n=== {name} n={n} ===")
    try:
        out = fn(n)
    except Exception as exc:
        print(f"RAISED {type(exc).__name__}: {exc}")
        return False
    row = np.asarray(out["row"], dtype=np.int64)
    col = np.asarray(out["col"], dtype=np.int64)
    w = np.asarray(out["weight"], dtype=np.float64)
    exp_row = np.arange(n, dtype=np.int64)
    exp_col = (np.arange(n, dtype=np.int64) * 10)
    exp_w = np.arange(n, dtype=np.float64) * 0.5
    ok = True
    ok &= check(f"{name} row", row, exp_row)
    ok &= check(f"{name} col", col, exp_col)
    ok &= check(f"{name} weight", w, exp_w)
    return ok


def main() -> int:
    print(f"numpy {np.__version__}")
    try:
        import subprocess

        print(subprocess.getoutput("mojo --version").strip())
    except Exception as exc:
        print("mojo --version failed:", exc)

    all_ok = True
    for n in (1, 2, 8, 64, 4096):
        exp_f = np.arange(n, dtype=np.float64)
        exp_i = np.arange(n, dtype=np.int64)
        for label, fn, exp in (
            ("fill_f64", ker.fill_f64, exp_f),
            ("fill_i64", ker.fill_i64, exp_i),
            ("fill_f64_parallel", ker.fill_f64_parallel, exp_f),
        ):
            try:
                all_ok &= check(f"{label} n={n}", np.asarray(fn(n)), exp)
            except Exception as exc:
                print(f"\n=== {label} n={n} RAISED ===\n{type(exc).__name__}: {exc}")
                all_ok = False

        all_ok &= check_stencil(f"stencil_like", n, ker.stencil_like)
        all_ok &= check_stencil(f"stencil_like_parallel", n, ker.stencil_like_parallel)

    try:
        all_ok &= check(
            "fill_f64_via_python_index n=8",
            np.asarray(ker.fill_f64_via_python_index(8)),
            np.arange(8, dtype=np.float64),
        )
    except Exception as exc:
        print(f"\n=== python index RAISED ===\n{type(exc).__name__}: {exc}")
        all_ok = False

    print("\n" + ("ALL OK" if all_ok else "FAILED (see mismatches / raises above)"))
    print(
        "\nIf ALL OK on your build: the historical aether bug may need a larger "
        "COO compact / SEGV case — ask Modular what the supported numpy write "
        "API is anyway (list append is too slow for production)."
    )
    return 0 if all_ok else 1


if __name__ == "__main__":
    raise SystemExit(main())

Questions for Modular

Is UnsafePointer(unsafe_from_address=arr.ctypes.data) the supported way to mutate an existing NumPy ndarray from Mojo?
Are there any known issues with multiple live pointers into different NumPy buffers in the same function or under parallelize?
What is the recommended high-performance alternative to Python lists + np.asarray() for returning large float64 buffers to Python?

Curious · August 5, 2026, 5:12pm

Conclusion

The bug was a read-after-free of the source array, not corrupt stores.

UnsafePointer(unsafe_from_address=Int(py=arr.ctypes.data)) has MutUntrackedOrigin, so it does not keep the numpy object alive; the array was destroyed after the pointer was taken and the kernel then read freed memory. Weights looked like heap addresses because they were computed from a freed fx, while pure-store channels (row, col) always looked fine.

Every case that reads a numpy buffer through a raw pointer fails at n ≥ 4096 (read_one_write_f64, stencil_like*, stencil_like_list_fx, stencil_like_bitcast). Cases that keep the object alive pass: stencil_like_keepalive (touches the objects after the loop) and stencil_like_reacquire (re-takes the pointer each iteration). Pure-store cases never fail.

The fix

std.python.numpy.from_numpy_array — the span’s origin is tied to the numpy binding, so the compiler keeps the array alive:

from std.python.numpy import from_numpy_array

var fx_s = from_numpy_arrayfx
var w_s = from_numpy_arrayweight

w_s\[i\] = fx_s\[i\] \* 0.5

Mutability follows the binding, so rebind a function argument to a local var before borrowing it for writes. The span_* and option_c_* cases pass at every size, serial and under parallelize.

The failing raw-pointer cases are kept deliberately as regression coverage.

AlisterBurt · August 7, 2026, 12:28am

Hi @Curious, I’ve hit similar issues in the past because of the ASAP destruction, definitely a gotcha. The correct place for reporting issues like this one would be a github issue over on the modular repo, would you like to open one? From there it’s more likely to be seen by engineering internally and lead to change

Curious · August 7, 2026, 1:00am

Well this was embarassing that I did not catch it and is not a bug on your end like the sm90a bug but just a missuse on my end. I do not want to polute the issues with dumb questions (already posted probably the dumbest question ever on there).

Topic		Replies	Views
High performance, fixed size 1D and 2D arrays on a CPU General	4	239	July 29, 2025
Matrix struct with Buffer: error at creating Standard Library discussion , mojo-compiler	0	62	April 17, 2026
How to convert numpy array items to mojo Float? Mojo discussion , 24_6	4	266	February 5, 2025
Mojo manual gpu basics exercise does not compile GPU Programming 25_3	7	289	April 2, 2025
Python.array() or Python.memoryview() converters for Mojo types Mojo gpu , python-interop	0	90	August 11, 2025

Writing Float64 into numpy via UnsafePointer corrupts when also reading another float64 buffer (Mojo 1.0b3)

Conclusion

The fix

Related topics