File size: 8,744 Bytes
f7ccd08
 
 
32b4893
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
// SPDX-License-Identifier: LicenseRef-Sovereign-Leviathan-SL-AGPL3-001
// Licensing: see LICENSE in the repository root.

//! Memory address calculation and stride handling
//!
//! Core address formula:
//!   addr = base + offset + Σ(index[i] * stride[i])

use crate::error::{HardwareError, Result};

/// Memory domain for address allocation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryDomain {
    /// Standard heap allocation
    Heap,
    /// GPU/accelerator memory (future)
    Accelerator,
    /// Pinned memory for DMA
    Pinned,
}

/// Address calculation context
#[derive(Debug, Clone)]
pub struct MemoryLayout {
    /// Physical buffer capacity
    pub capacity: usize,

    /// Base offset for this tensor in buffer
    pub offset: usize,

    /// Shape dimensions
    pub shape: Vec<usize>,

    /// Stride for each dimension
    pub strides: Vec<usize>,

    /// Memory domain
    pub domain: MemoryDomain,
}

impl MemoryLayout {
    /// Create new memory layout
    ///
    /// **Validation**: offset must be in bounds, shape must not exceed capacity
    pub fn new(
        capacity: usize,
        offset: usize,
        shape: Vec<usize>,
        strides: Vec<usize>,
    ) -> Result<Self> {
        if shape.is_empty() {
            return Err(HardwareError::ZeroElementShape);
        }

        if strides.len() != shape.len() {
            return Err(HardwareError::RankMismatch {
                expected: shape.len(),
                actual: strides.len(),
            });
        }

        // Check if any dimension is zero
        if shape.iter().any(|&s| s == 0) {
            return Err(HardwareError::ZeroElementShape);
        }

        // Validate offset is within capacity
        if offset >= capacity {
            return Err(HardwareError::AddressOutOfBounds {
                address: offset,
                capacity,
            });
        }

        let layout = MemoryLayout {
            capacity,
            offset,
            shape,
            strides,
            domain: MemoryDomain::Heap,
        };

        Ok(layout)
    }

    /// Calculate total number of elements
    #[inline]
    pub fn element_count(&self) -> usize {
        self.shape.iter().product()
    }

    /// Calculate address for multi-dimensional index
    ///
    /// **Formula**: addr = offset + Σ(index[i] * stride[i])
    pub fn address_for(&self, indices: &[usize]) -> Result<usize> {
        if indices.len() != self.shape.len() {
            return Err(HardwareError::RankMismatch {
                expected: self.shape.len(),
                actual: indices.len(),
            });
        }

        // Bounds check each index
        for (i, &idx) in indices.iter().enumerate() {
            if idx >= self.shape[i] {
                return Err(HardwareError::IndexOutOfBounds {
                    index: idx,
                    size: self.shape[i],
                });
            }
        }

        // Calculate strided address
        let mut addr = self.offset;
        for (i, &idx) in indices.iter().enumerate() {
            addr = addr.saturating_add(idx.saturating_mul(self.strides[i]));
        }

        // Validate address is within capacity
        if addr >= self.capacity {
            return Err(HardwareError::AddressOutOfBounds {
                address: addr,
                capacity: self.capacity,
            });
        }

        Ok(addr)
    }

    /// Check if strides are canonical (row-major)
    ///
    /// Canonical: stride[i] = ∏_{j>i} shape[j]
    pub fn is_canonical_strides(&self) -> bool {
        if self.shape.is_empty() {
            return false;
        }

        let mut expected_stride = 1;
        for i in (0..self.shape.len()).rev() {
            if self.strides[i] != expected_stride {
                return false;
            }
            expected_stride = expected_stride.saturating_mul(self.shape[i]);
        }
        true
    }

    /// Compute canonical strides for current shape
    pub fn canonical_strides(&self) -> Vec<usize> {
        let mut strides = vec![0; self.shape.len()];
        let mut stride: usize = 1;
        for i in (0..self.shape.len()).rev() {
            strides[i] = stride;
            stride = stride.saturating_mul(self.shape[i]);
        }
        strides
    }

    /// Create a view (slice) of this layout
    ///
    /// **Arguments**: ranges = [(start, end, step), ...]
    /// **Returns**: new layout with updated shape, strides, offset
    pub fn slice(&self, ranges: &[(usize, usize, usize)]) -> Result<Self> {
        if ranges.len() != self.shape.len() {
            return Err(HardwareError::RankMismatch {
                expected: self.shape.len(),
                actual: ranges.len(),
            });
        }

        let mut new_shape = Vec::with_capacity(self.shape.len());
        let mut new_strides = Vec::with_capacity(self.strides.len());
        let mut new_offset = self.offset;

        for (dim, &(start, end, step)) in ranges.iter().enumerate() {
            // Validate range
            if start > end || end > self.shape[dim] || step == 0 {
                return Err(HardwareError::IndexOutOfBounds {
                    index: start,
                    size: self.shape[dim],
                });
            }

            let new_size = (end - start + step - 1) / step; // Ceiling division
            new_shape.push(new_size);

            // Update offset for this dimension
            new_offset = new_offset.saturating_add(start.saturating_mul(self.strides[dim]));

            // Update stride (multiply by step)
            new_strides.push(self.strides[dim].saturating_mul(step));
        }

        MemoryLayout {
            capacity: self.capacity,
            offset: new_offset,
            shape: new_shape,
            strides: new_strides,
            domain: self.domain,
        }
        .validate_bounds()
    }

    /// Validate that maximum address is within capacity
    fn validate_bounds(&self) -> Result<Self> {
        if self.element_count() == 0 {
            return Err(HardwareError::ZeroElementShape);
        }

        // Calculate worst-case address (last element)
        let mut max_indices = vec![0; self.shape.len()];
        for (i, &s) in self.shape.iter().enumerate() {
            if s > 0 {
                max_indices[i] = s - 1;
            }
        }

        let _ = self.address_for(&max_indices)?;
        Ok(self.clone())
    }
}

/// Flat (1D) linear address indexing
#[derive(Debug, Clone)]
pub struct LinearIndex {
    pub capacity: usize,
    pub offset: usize,
}

impl LinearIndex {
    pub fn new(capacity: usize, offset: usize) -> Result<Self> {
        if offset >= capacity {
            return Err(HardwareError::AddressOutOfBounds {
                address: offset,
                capacity,
            });
        }
        Ok(LinearIndex { capacity, offset })
    }

    /// Get address for linear index
    #[inline]
    pub fn address(&self, index: usize) -> Result<usize> {
        let addr = self.offset.saturating_add(index);
        if addr >= self.capacity {
            return Err(HardwareError::AddressOutOfBounds {
                address: addr,
                capacity: self.capacity,
            });
        }
        Ok(addr)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_memory_layout_creation() {
        let layout = MemoryLayout::new(1000, 0, vec![10, 10], vec![10, 1]).unwrap();
        assert_eq!(layout.element_count(), 100);
    }

    #[test]
    fn test_address_calculation() {
        let layout = MemoryLayout::new(1000, 0, vec![10, 10], vec![10, 1]).unwrap();
        let addr = layout.address_for(&[2, 3]).unwrap();
        assert_eq!(addr, 2 * 10 + 3 * 1); // 23
    }

    #[test]
    fn test_canonical_strides_check() {
        let layout = MemoryLayout::new(1000, 0, vec![10, 10], vec![10, 1]).unwrap();
        assert!(layout.is_canonical_strides());

        let layout2 = MemoryLayout::new(1000, 0, vec![10, 10], vec![5, 1]).unwrap();
        assert!(!layout2.is_canonical_strides());
    }

    #[test]
    fn test_bounds_checking() {
        let layout = MemoryLayout::new(100, 0, vec![10, 10], vec![10, 1]).unwrap();
        assert!(layout.address_for(&[5, 5]).is_ok());
        assert!(layout.address_for(&[10, 5]).is_err());
    }

    #[test]
    fn test_slice_operation() {
        let layout = MemoryLayout::new(1000, 0, vec![10, 10], vec![10, 1]).unwrap();
        let sliced = layout.slice(&[(2, 5, 1), (3, 7, 1)]).unwrap();
        assert_eq!(sliced.shape, vec![3, 4]);
    }

    #[test]
    fn test_zero_element_shape_rejected() {
        assert!(MemoryLayout::new(1000, 0, vec![10, 0], vec![10, 1]).is_err());
    }
}