burn-mlx 0.1.2

MLX backend for Burn deep learning framework - native Apple Silicon GPU acceleration
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

//! MLX Backend implementation for Burn.

use burn_tensor::backend::Backend;
use burn_tensor::TensorMetadata;
use mlx_rs::Array;
use std::sync::atomic::{AtomicU64, Ordering};

use crate::device::MlxDevice;

// Global seed for random number generation
static SEED: AtomicU64 = AtomicU64::new(0);

/// MLX tensor primitive with shape tracking.
#[derive(Debug, Clone)]
pub struct MlxTensorPrimitive {
    /// The underlying MLX array.
    pub array: Array,
    /// Cached shape for fast access.
    pub shape: Vec<usize>,
}

impl MlxTensorPrimitive {
    /// Create a new tensor primitive.
    pub fn new(array: Array) -> Self {
        let shape = array.shape().iter().map(|&s| s as usize).collect();
        Self { array, shape }
    }

    /// Get the array reference.
    pub fn array(&self) -> &Array {
        &self.array
    }

    /// Get the shape.
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
}

// SAFETY: MLX arrays can be sent between threads.
// MLX uses internal synchronization for its compute graph.
unsafe impl Send for MlxTensorPrimitive {}
unsafe impl Sync for MlxTensorPrimitive {}

impl TensorMetadata for MlxTensorPrimitive {
    fn dtype(&self) -> burn_tensor::DType {
        // Map MLX dtype to Burn dtype
        match self.array.dtype() {
            mlx_rs::Dtype::Float32 => burn_tensor::DType::F32,
            mlx_rs::Dtype::Float16 => burn_tensor::DType::F16,
            mlx_rs::Dtype::Bfloat16 => burn_tensor::DType::BF16,
            mlx_rs::Dtype::Float64 => burn_tensor::DType::F64,
            mlx_rs::Dtype::Int32 => burn_tensor::DType::I32,
            mlx_rs::Dtype::Int64 => burn_tensor::DType::I64,
            mlx_rs::Dtype::Bool => burn_tensor::DType::Bool,
            _ => burn_tensor::DType::F32, // Default fallback
        }
    }

    fn shape(&self) -> burn_tensor::Shape {
        burn_tensor::Shape::from(self.shape.clone())
    }
}

/// Quantized tensor primitive (placeholder for future implementation).
#[derive(Debug, Clone)]
pub struct MlxQuantizedTensorPrimitive {
    /// The underlying tensor (stored as float for now).
    pub tensor: MlxTensorPrimitive,
    /// Quantization scheme.
    pub scheme: QuantizationScheme,
}

/// Quantization scheme.
#[derive(Debug, Clone, Copy, Default)]
pub enum QuantizationScheme {
    #[default]
    None,
}

// SAFETY: Same as MlxTensorPrimitive
unsafe impl Send for MlxQuantizedTensorPrimitive {}
unsafe impl Sync for MlxQuantizedTensorPrimitive {}

impl TensorMetadata for MlxQuantizedTensorPrimitive {
    fn dtype(&self) -> burn_tensor::DType {
        self.tensor.dtype()
    }

    fn shape(&self) -> burn_tensor::Shape {
        burn_tensor::Shape::from(self.tensor.shape.clone())
    }
}

impl burn_tensor::quantization::QTensorPrimitive for MlxQuantizedTensorPrimitive {
    fn scheme(&self) -> &burn_tensor::quantization::QuantizationScheme {
        // Return a reference to a static scheme
        static SYMMETRIC: burn_tensor::quantization::QuantizationScheme =
            burn_tensor::quantization::QuantizationScheme::PerTensorSymmetric(
                burn_tensor::quantization::QuantizationType::QInt8,
            );
        &SYMMETRIC
    }
}

/// MLX Backend for Burn.
#[derive(Debug, Default, Clone, Copy)]
pub struct Mlx;

impl Backend for Mlx {
    type Device = MlxDevice;

    type FloatTensorPrimitive = MlxTensorPrimitive;
    type FloatElem = f32;

    type IntTensorPrimitive = MlxTensorPrimitive;
    type IntElem = i32;

    type BoolTensorPrimitive = MlxTensorPrimitive;
    type BoolElem = bool;

    type QuantizedTensorPrimitive = MlxQuantizedTensorPrimitive;
    type QuantizedEncoding = i8;

    fn name() -> String {
        "mlx".to_string()
    }

    fn seed(seed: u64) {
        SEED.store(seed, Ordering::SeqCst);
        // MLX uses its own seeding mechanism
        mlx_rs::random::seed(seed);
    }

    fn sync(device: &Self::Device) {
        // MLX is lazy-evaluated; sync forces evaluation
        // This is a no-op in MLX as synchronization happens implicitly
        // when reading tensor values
        let _ = device;
    }
}

/// Get the current seed value.
pub fn get_seed() -> u64 {
    SEED.load(Ordering::SeqCst)
}