[core][GPU][StaticMatrix] Introduce GPU backend for NuMojo!

[shivasankarka]

This PR introduces initial GPU support for Numojo #273

It adds unified device and storage abstractions, a basic matrix representation for GPU computations, and several core GPU kernels (elementwise add/sub/mul, matmul, fill, and a block-level reduction). This work lays the foundation for using Mojo GPU features to accelerate array operations.

The design is inspired by PyTorch Tensor while keeping NumPy-like API choices where possible.

Notes

• The StaticMatrix is still a very basic structure with only some getter and setter functions to showcase the proof of concept of a GPU backend in NuMojo. We will expand in future to include all features from Matrix type.
• It's named as StaticMatrix as a compile time shape and strides would help optimize a lot of the loops and gpu kernels! This would be a Matrix type that takes advantage of Mojo's compile time capabilities as much as possible! We will modify the API to support compile time optimisations in future updates.

---

What’s Included

Device & context abstraction

• numojo/core/gpu/device.mojo — device and context primitives to target GPU

Unified storage

• numojo/core/gpu/storage.mojo — unified CPU/GPU memory management for buffers

Matrix primitives

• numojo/core/staticmatrix.mojo — adds a StaticMatrix struct to prototype GPU usage before extending to N-D arrays

GPU kernels

• numojo/core/gpu/matrix_kernels.mojo — implements:
  • Vectorized elementwise kernels: add, mul, fill (and sub)
  • Tiled matmul helpers
  • matrix_reduce_sum_kernel (per-block reduction)

Other

• Launch-parameter helpers and dtype-specialization hooks for future optimizations
• Updated for latest Mojo nightly (pixi updates)
• Small dtype fixes and deprecation error fixes

---

Example

fn main() raises:
    alias SIZE: Int = 1024
    alias cpu: Device = Device.CPU
    alias mps: Device = Device.MPS

    var arr_cpu_1 = StaticMatrix[DType.float32](shape=(SIZE, SIZE), order="C", fill_value=1.0)
    var arr_cpu_2 = StaticMatrix[DType.float32]((SIZE, SIZE), fill_value=2.0)
    var matmul_cpu = arr_cpu_1 @ arr_cpu_2
    print(matmul_cpu)

    var arr_gpu_1 = StaticMatrix[DType.float32, device=mps](shape=(SIZE, SIZE), order="C", fill_value=1.0)
    var arr_gpu_2 = StaticMatrix[DType.float32, device=mps](shape=(SIZE, SIZE), order="C", fill_value=2.0)
    var matmul_gpu = arr_gpu_1 @ arr_gpu_2
    print(matmul_gpu)

    var arr_gpu_fromcpu_1 = arr_cpu_1.to[mps]()
    var arr_gpu_fromcpu_2 = arr_cpu_2.to[mps]()
    var matmul_gpu_fromcpu = arr_gpu_fromcpu_1 @ arr_gpu_fromcpu_2
    print(matmul_gpu_fromcpu)