NDArray implementation in mojo

最新推荐文章于 2026-06-19 17:00:49 发布

原创最新推荐文章于 2026-06-19 17:00:49 发布 · 376 阅读

0 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#mojo #python #开发语言

文章讨论了在没有继承支持的情况下，如何通过复合方法扩展NDBuffer来创建NDArray类，用于实现AI功能。NDArray结构体包含了NDBuffer，并实现了__getitem__和__setitem__方法以支持SIMD操作。由于不支持关键字参数，因此无法直接使用如x[0,0]的形式访问元素，而需使用StaticIntTuple。

To develop AI function we need a ndarray like class. But mojo does not have this kind of class now. So we need to
implement one. The most nature way is to extend NDBuffer in mojo. But mojo does not support inheritance now. In this
article I will use composite instead of inheritance to extend NDBuffer.
First we need import necessary libs:

# common import
from String import String
from Bool import Bool
from List import VariadicList
from Buffer import NDBuffer
from List import DimList
from DType import DType
from Pointer import DTypePointer
from TargetInfo import dtype_sizeof, dtype_simd_width
from Index import StaticIntTuple

alias nelts = dtype_simd_width[DType.f32]()

Now it is time to extend NDBuffer:

struct NDArray[rank:Int, dims:DimList, dtype:DType]:
    var ndb:NDBuffer[rank, dims, dtype]
    
    fn __init__(inout self, size:Int):
        let data = DTypePointer[dtype].alloc(size)
        self.ndb = NDBuffer[rank, dims, dtype](data)
        
    fn __getitem__(self, idxs:StaticIntTuple[rank]) -> SIMD[dtype,1]:
        return self.ndb.simd_load[1](idxs)
    
    fn __setitem__(self, idxs:StaticIntTuple[rank], val:SIMD[dtype,1]):
        self.ndb.simd_store[1](idxs, val)

To use the NDArray:

alias x_rank = 2
var x = NDArray[x_rank, DimList(2, 3), DType.f32](6)
x[StaticIntTuple[x_rank](1,1)] = 999.9
print(x[StaticIntTuple[x_rank](0,0)], x[StaticIntTuple[x_rank](1,1)])

Discussion

As see the above code you may be ask why not use x[0,0] to get and set the item of NDArray. The reason is that mojo
dose not support named keyword arguments now. So if I define the NDArray like these:

struct ndarray[rank:Int, dims:DimList, dtype:DType]:
    var ndb:NDBuffer[rank, dims, dtype]
    
    fn __init__(inout self, size:Int):
        let data = DTypePointer[dtype].alloc(size)
        self.ndb = NDBuffer[rank, dims, dtype](data)
        
    fn __getitem__(self, *idxs:Int) -> SIMD[dtype,1]:
        return self.ndb.simd_load[1](idxs)
    
    fn __setitem__(self, *idxs:Int, val:SIMD[dtype,1]):
        self.ndb.simd_store[1](idxs, val)

When I access x[0,0] it will report these errors:

error: Expression [23]:16:37: TODO: keyword arguments not supported yet
    fn __setitem__(self, *idxs:Int, val:SIMD[dtype,1]):
                                    ^

error: Expression [23]:17:31: no matching function in call to 'simd_store': 
        self.ndb.simd_store[1](idxs, val)
        ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~

/.modular/Kernels/mojo/Stdlib/Buffer.mojo:1173:5: candidate not viable: method argument #0 cannot be converted from 'variadic<_"$Int"::_Int>' to 'StaticIntTuple[rank]'
    fn simd_store[
    ^

/.modular/Kernels/mojo/Stdlib/Buffer.mojo:1191:5: candidate not viable: method argument #0 cannot be converted from 'variadic<_"$Int"::_Int>' to 'StaticTuple[rank, Int]'
    fn simd_store[
    ^