value.rs 6.66 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

20
use std::{os::raw::c_char, str::FromStr};
21

22
use crate::ffi::*;
23

24
impl DLDataType {
25 26 27 28 29
    fn new(type_code: u8, bits: u8, lanes: u16) -> Self {
        Self {
            code: type_code,
            bits,
            lanes,
30 31 32 33
        }
    }
}

34 35 36 37 38 39 40 41
#[derive(Debug, Fail)]
pub enum ParseTvmTypeError {
    #[fail(display = "invalid number: {}", _0)]
    InvalidNumber(std::num::ParseIntError),
    #[fail(display = "unknown type: {}", _0)]
    UnknownType(String),
}

42 43
/// Implements TVMType conversion from `&str` of general format `{dtype}{bits}x{lanes}`
/// such as "int32", "float32" or with lane "float32x1".
44 45
impl FromStr for DLDataType {
    type Err = ParseTvmTypeError;
46 47
    fn from_str(type_str: &str) -> Result<Self, Self::Err> {
        if type_str == "bool" {
48
            return Ok(DLDataType::new(1, 1, 1));
49 50
        }

51
        let mut type_lanes = type_str.split('x');
52 53 54 55
        let typ = type_lanes.next().expect("Missing dtype");
        let lanes = type_lanes
            .next()
            .map(|l| <u16>::from_str_radix(l, 10))
56 57
            .unwrap_or(Ok(1))
            .map_err(ParseTvmTypeError::InvalidNumber)?;
58 59 60
        let (type_name, bits) = match typ.find(char::is_numeric) {
            Some(idx) => {
                let (name, bits_str) = typ.split_at(idx);
61 62 63 64
                (
                    name,
                    u8::from_str_radix(bits_str, 10).map_err(ParseTvmTypeError::InvalidNumber)?,
                )
65
            }
66
            None => (typ, 32),
67 68
        };

69 70 71 72 73
        let type_code = match type_name {
            "int" => 0,
            "uint" => 1,
            "float" => 2,
            "handle" => 3,
74
            _ => return Err(ParseTvmTypeError::UnknownType(type_name.to_string())),
75 76
        };

77
        Ok(DLDataType::new(type_code, bits, lanes))
78 79 80
    }
}

81
impl std::fmt::Display for DLDataType {
82 83 84
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        if self.bits == 1 && self.lanes == 1 {
            return write!(f, "bool");
85
        }
86 87 88 89 90 91
        let mut type_str = match self.code {
            0 => "int",
            1 => "uint",
            2 => "float",
            4 => "handle",
            _ => "unknown",
92
        }
93
        .to_string();
94

95 96 97
        type_str += &self.bits.to_string();
        if self.lanes > 1 {
            type_str += &format!("x{}", self.lanes);
98
        }
99
        f.write_str(&type_str)
100 101 102
    }
}

103 104 105 106 107 108
macro_rules! impl_pod_tvm_value {
    ($field:ident, $field_ty:ty, $( $ty:ty ),+) => {
        $(
            impl From<$ty> for TVMValue {
                fn from(val: $ty) -> Self {
                    TVMValue { $field: val as $field_ty }
109 110 111
                }
            }

112 113 114
            impl From<TVMValue> for $ty {
                fn from(val: TVMValue) -> Self {
                    unsafe { val.$field as $ty }
115 116
                }
            }
117
        )+
118
    };
119 120 121
    ($field:ident, $ty:ty) => {
        impl_pod_tvm_value!($field, $ty, $ty);
    }
122 123
}

124 125
impl_pod_tvm_value!(v_int64, i64, i8, u8, i16, u16, i32, u32, i64, u64, isize, usize);
impl_pod_tvm_value!(v_float64, f64, f32, f64);
126
impl_pod_tvm_value!(v_type, DLDataType);
127
impl_pod_tvm_value!(v_ctx, TVMContext);
128

129 130 131 132
#[derive(Debug, Fail)]
#[fail(display = "unsupported device: {}", _0)]
pub struct UnsupportedDeviceError(String);

133 134 135 136
macro_rules! impl_tvm_context {
    ( $( $dev_type:ident : [ $( $dev_name:ident ),+ ] ),+ ) => {
        /// Creates a TVMContext from a string (e.g., "cpu", "gpu", "ext_dev")
        impl FromStr for TVMContext {
137
            type Err = UnsupportedDeviceError;
138 139 140 141
            fn from_str(type_str: &str) -> Result<Self, Self::Err> {
                Ok(Self {
                    device_type: match type_str {
                         $( $(  stringify!($dev_name)  )|+ => $dev_type ),+,
142
                        _ => return Err(UnsupportedDeviceError(type_str.to_string())),
143 144 145
                    },
                    device_id: 0,
                })
146 147 148
            }
        }

149 150
        impl TVMContext {
            $(
151
                $(
152 153 154 155 156 157
                    pub fn $dev_name(device_id: usize) -> Self {
                        Self {
                            device_type: $dev_type,
                            device_id: device_id as i32,
                        }
                    }
158
                )+
159
            )+
160
        }
161
    };
162
}
163 164 165 166 167 168 169 170 171 172

impl_tvm_context!(
    DLDeviceType_kDLCPU: [cpu, llvm, stackvm],
    DLDeviceType_kDLGPU: [gpu, cuda, nvptx],
    DLDeviceType_kDLOpenCL: [cl],
    DLDeviceType_kDLMetal: [metal],
    DLDeviceType_kDLVPI: [vpi],
    DLDeviceType_kDLROCM: [rocm],
    DLDeviceType_kDLExtDev: [ext_dev]
);
173

174 175 176 177 178 179
/// A struct holding TVM byte-array.
///
/// ## Example
///
/// ```
/// let v = b"hello";
180
/// let barr = tvm_common::TVMByteArray::from(&v);
181 182 183
/// assert_eq!(barr.len(), v.len());
/// assert_eq!(barr.data(), &[104u8, 101, 108, 108, 111]);
/// ```
184
impl TVMByteArray {
185
    /// Gets the underlying byte-array
186 187 188
    pub fn data(&self) -> &'static [u8] {
        unsafe { std::slice::from_raw_parts(self.data as *const u8, self.size) }
    }
189 190 191 192 193 194 195 196 197 198

    /// Gets the length of the underlying byte-array
    pub fn len(&self) -> usize {
        self.size
    }

    /// Converts the underlying byte-array to `Vec<u8>`
    pub fn to_vec(&self) -> Vec<u8> {
        self.data().to_vec()
    }
199 200 201 202

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
203 204
}

205 206 207 208 209 210 211
// Needs AsRef for Vec
impl<T: AsRef<[u8]>> From<T> for TVMByteArray {
    fn from(arg: T) -> Self {
        let arg = arg.as_ref();
        TVMByteArray {
            data: arg.as_ptr() as *const c_char,
            size: arg.len(),
212 213 214
        }
    }
}
215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231

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

    #[test]
    fn convert() {
        let v = vec![1u8, 2, 3];
        let barr = TVMByteArray::from(&v);
        assert_eq!(barr.len(), v.len());
        assert_eq!(barr.to_vec(), vec![1u8, 2, 3]);
        let v = b"hello";
        let barr = TVMByteArray::from(&v);
        assert_eq!(barr.len(), v.len());
        assert_eq!(barr.data(), &[104u8, 101, 108, 108, 111]);
    }
}