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

use crate::{
    libs::cache::{Cache, CacheReturn}, AsDev, BaseDevice,
    Buffer, CDatatype, CacheBuf, ClearBuf, Alloc, GenericBlas, ManualMem, VecRead, WriteBuf, Device, DeviceType
};
use std::{ffi::c_void, fmt::Debug, alloc::{Layout, handle_alloc_error}, mem::size_of, cell::RefCell};

use super::RawCpuBuf;

#[derive(Debug, Default)]
/// A CPU is used to perform calculations on the host CPU.
/// To make new operations invocable, a trait providing new functions should be implemented for [CPU].
///
/// # Example
/// ```
/// use custos::{CPU, VecRead, Buffer};
///
/// let device = CPU::new();
/// let a = Buffer::from((&device, [1, 2, 3]));
///
/// let out = device.read(&a);
///
/// assert_eq!(out, vec![1, 2, 3]);
/// ```
pub struct CPU {
    pub cache: RefCell<Cache<RawCpuBuf>>
}

impl CPU {
    /// Creates an [CPU] with an InternCPU that holds an empty vector of pointers.
    #[must_use]
    pub fn new() -> CPU {
        CPU {
            cache: RefCell::new(Cache::default()),
        }
    }
}

impl<T: Clone + Default> Alloc<T> for CPU {
    fn alloc(&self, len: usize) -> (*mut T, *mut c_void, u64) {
        assert!(len > 0, "invalid buffer len: 0");
        let layout = Layout::array::<T>(len).unwrap();
        let ptr = unsafe {
            std::alloc::alloc(layout)
        };

        // initialize block of memory
        for element in unsafe {std::slice::from_raw_parts_mut(ptr, len*size_of::<T>())} {
            *element = 0;
        }

        if ptr.is_null() {
            handle_alloc_error(layout);
        }
        (ptr as *mut T, std::ptr::null_mut(), 0)
    }

    fn with_data(&self, data: &[T]) -> (*mut T, *mut c_void, u64) {
        assert!(!data.is_empty(), "invalid buffer len: 0");
        let (ptr, _, _) = self.alloc(data.len());
        let slice = unsafe {std::slice::from_raw_parts_mut(ptr, data.len())};
        slice.clone_from_slice(data);
        (ptr, std::ptr::null_mut(), 0)
    }
    fn alloc_with_vec(&self, mut vec: Vec<T>) -> (*mut T, *mut c_void, u64) {
        assert!(!vec.is_empty(), "invalid buffer len: 0");
    
        let ptr = vec.as_mut_ptr();
        std::mem::forget(vec);
        (ptr, std::ptr::null_mut(), 0)
    }

    fn as_dev(&self) -> crate::Device {
        Device {
            device_type: DeviceType::CPU,
            device: self as *const CPU as *mut u8
        }
    }
}

impl AsDev for CPU {}

impl<T> ManualMem<T> for CPU {
    fn drop_buf(&self, buf: Buffer<T>) {
        unsafe {
            drop(Box::from_raw(buf.ptr.0));
        }
    }
}

impl CacheReturn<RawCpuBuf> for CPU {
    #[inline]
    fn cache(&self) -> std::cell::RefMut<Cache<RawCpuBuf>> {
        self.cache.borrow_mut()
    }
}

impl<'a, T: Copy + Default> CacheBuf<'a, T> for CPU {
    #[inline]
    fn cached(&'a self, len: usize) -> Buffer<'a, T> {
        Cache::get::<T, CPU>(self, len)
    }
}

#[inline]
pub fn cpu_cached<T: Copy+Default>(device: &CPU, len: usize) -> Buffer<T> {
    device.cached(len)
}

impl<T: Copy + Default> VecRead<T> for CPU {
    fn read(&self, buf: &Buffer<T>) -> Vec<T> {
        buf.as_slice().to_vec()
    }
}

impl<T: Default> ClearBuf<T> for CPU {
    fn clear(&self, buf: &mut Buffer<T>) {
        for value in buf {
            *value = T::default();
        }
    }
}

impl<T: Copy> WriteBuf<T> for CPU {
    fn write(&self, buf: &mut Buffer<T>, data: &[T]) {
        buf.copy_from_slice(data)
    }
}

impl<T: CDatatype + GenericBlas> BaseDevice<T> for CPU {}