perf(flex): reduce int cast codegen size (#5357)
* ignore additional files created by profiling * reduce absurdly large cast_int codegen
This commit is contained in:
laund
@@ -25,3 +25,10 @@ uv.lock
|
||||
tags
|
||||
|
||||
examples/**/Cargo.lock
|
||||
|
||||
# rustc self-profiling files
|
||||
*.mm_profdata
|
||||
*.stacks_folded
|
||||
*profiler.json
|
||||
# linker maps
|
||||
*.map
|
||||
|
||||
@@ -11,6 +11,7 @@ extend-exclude = [
|
||||
"*.onnx",
|
||||
"*.proto",
|
||||
"assets/ModuleSerialization.xml",
|
||||
"*.map", # linker map files
|
||||
]
|
||||
|
||||
[default.extend-words]
|
||||
|
||||
+119
-68
@@ -684,97 +684,148 @@ impl IntTensorOps<Flex> for Flex {
|
||||
|
||||
// Helper macro to convert between types
|
||||
macro_rules! cast_impl {
|
||||
($src_type:ty, $dst_type:ty, $dst_dtype:expr) => {{
|
||||
let src: &[$src_type] = tensor.storage();
|
||||
let dst: Vec<$dst_type> = src.iter().map(|&x| x as $dst_type).collect();
|
||||
FlexTensor::new(
|
||||
Bytes::from_elems(dst),
|
||||
Layout::contiguous(shape),
|
||||
$dst_dtype,
|
||||
)
|
||||
($storage:ident, $dst_type:ty) => {{
|
||||
Some(Bytes::from_elems(
|
||||
$storage
|
||||
.iter()
|
||||
.map(|&x| x as $dst_type)
|
||||
.collect::<Vec<$dst_type>>(),
|
||||
))
|
||||
}};
|
||||
}
|
||||
|
||||
// Match source dtype to target dtype
|
||||
match (tensor.dtype(), target_dtype) {
|
||||
let bytes = match tensor.dtype() {
|
||||
// From I64
|
||||
(DType::I64, DType::I32) => cast_impl!(i64, i32, DType::I32),
|
||||
(DType::I64, DType::I16) => cast_impl!(i64, i16, DType::I16),
|
||||
(DType::I64, DType::I8) => cast_impl!(i64, i8, DType::I8),
|
||||
(DType::I64, DType::U64) => cast_impl!(i64, u64, DType::U64),
|
||||
(DType::I64, DType::U32) => cast_impl!(i64, u32, DType::U32),
|
||||
(DType::I64, DType::U16) => cast_impl!(i64, u16, DType::U16),
|
||||
(DType::I64, DType::U8) => cast_impl!(i64, u8, DType::U8),
|
||||
DType::I64 => {
|
||||
let storage: &[i64] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From I32
|
||||
(DType::I32, DType::I64) => cast_impl!(i32, i64, DType::I64),
|
||||
(DType::I32, DType::I16) => cast_impl!(i32, i16, DType::I16),
|
||||
(DType::I32, DType::I8) => cast_impl!(i32, i8, DType::I8),
|
||||
(DType::I32, DType::U64) => cast_impl!(i32, u64, DType::U64),
|
||||
(DType::I32, DType::U32) => cast_impl!(i32, u32, DType::U32),
|
||||
(DType::I32, DType::U16) => cast_impl!(i32, u16, DType::U16),
|
||||
(DType::I32, DType::U8) => cast_impl!(i32, u8, DType::U8),
|
||||
DType::I32 => {
|
||||
let storage: &[i32] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From I16
|
||||
(DType::I16, DType::I64) => cast_impl!(i16, i64, DType::I64),
|
||||
(DType::I16, DType::I32) => cast_impl!(i16, i32, DType::I32),
|
||||
(DType::I16, DType::I8) => cast_impl!(i16, i8, DType::I8),
|
||||
(DType::I16, DType::U64) => cast_impl!(i16, u64, DType::U64),
|
||||
(DType::I16, DType::U32) => cast_impl!(i16, u32, DType::U32),
|
||||
(DType::I16, DType::U16) => cast_impl!(i16, u16, DType::U16),
|
||||
(DType::I16, DType::U8) => cast_impl!(i16, u8, DType::U8),
|
||||
DType::I16 => {
|
||||
let storage: &[i16] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From I8
|
||||
(DType::I8, DType::I64) => cast_impl!(i8, i64, DType::I64),
|
||||
(DType::I8, DType::I32) => cast_impl!(i8, i32, DType::I32),
|
||||
(DType::I8, DType::I16) => cast_impl!(i8, i16, DType::I16),
|
||||
(DType::I8, DType::U64) => cast_impl!(i8, u64, DType::U64),
|
||||
(DType::I8, DType::U32) => cast_impl!(i8, u32, DType::U32),
|
||||
(DType::I8, DType::U16) => cast_impl!(i8, u16, DType::U16),
|
||||
(DType::I8, DType::U8) => cast_impl!(i8, u8, DType::U8),
|
||||
DType::I8 => {
|
||||
let storage: &[i8] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From U64
|
||||
(DType::U64, DType::I64) => cast_impl!(u64, i64, DType::I64),
|
||||
(DType::U64, DType::I32) => cast_impl!(u64, i32, DType::I32),
|
||||
(DType::U64, DType::I16) => cast_impl!(u64, i16, DType::I16),
|
||||
(DType::U64, DType::I8) => cast_impl!(u64, i8, DType::I8),
|
||||
(DType::U64, DType::U32) => cast_impl!(u64, u32, DType::U32),
|
||||
(DType::U64, DType::U16) => cast_impl!(u64, u16, DType::U16),
|
||||
(DType::U64, DType::U8) => cast_impl!(u64, u8, DType::U8),
|
||||
DType::U64 => {
|
||||
let storage: &[u64] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From U32
|
||||
(DType::U32, DType::I64) => cast_impl!(u32, i64, DType::I64),
|
||||
(DType::U32, DType::I32) => cast_impl!(u32, i32, DType::I32),
|
||||
(DType::U32, DType::I16) => cast_impl!(u32, i16, DType::I16),
|
||||
(DType::U32, DType::I8) => cast_impl!(u32, i8, DType::I8),
|
||||
(DType::U32, DType::U64) => cast_impl!(u32, u64, DType::U64),
|
||||
(DType::U32, DType::U16) => cast_impl!(u32, u16, DType::U16),
|
||||
(DType::U32, DType::U8) => cast_impl!(u32, u8, DType::U8),
|
||||
DType::U32 => {
|
||||
let storage: &[u32] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From U16
|
||||
(DType::U16, DType::I64) => cast_impl!(u16, i64, DType::I64),
|
||||
(DType::U16, DType::I32) => cast_impl!(u16, i32, DType::I32),
|
||||
(DType::U16, DType::I16) => cast_impl!(u16, i16, DType::I16),
|
||||
(DType::U16, DType::I8) => cast_impl!(u16, i8, DType::I8),
|
||||
(DType::U16, DType::U64) => cast_impl!(u16, u64, DType::U64),
|
||||
(DType::U16, DType::U32) => cast_impl!(u16, u32, DType::U32),
|
||||
(DType::U16, DType::U8) => cast_impl!(u16, u8, DType::U8),
|
||||
DType::U16 => {
|
||||
let storage: &[u16] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U8 => cast_impl!(storage, u8),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// From U8
|
||||
(DType::U8, DType::I64) => cast_impl!(u8, i64, DType::I64),
|
||||
(DType::U8, DType::I32) => cast_impl!(u8, i32, DType::I32),
|
||||
(DType::U8, DType::I16) => cast_impl!(u8, i16, DType::I16),
|
||||
(DType::U8, DType::I8) => cast_impl!(u8, i8, DType::I8),
|
||||
(DType::U8, DType::U64) => cast_impl!(u8, u64, DType::U64),
|
||||
(DType::U8, DType::U32) => cast_impl!(u8, u32, DType::U32),
|
||||
(DType::U8, DType::U16) => cast_impl!(u8, u16, DType::U16),
|
||||
DType::U8 => {
|
||||
let storage: &[u8] = tensor.storage();
|
||||
match target_dtype {
|
||||
DType::I64 => cast_impl!(storage, i64),
|
||||
DType::I32 => cast_impl!(storage, i32),
|
||||
DType::I16 => cast_impl!(storage, i16),
|
||||
DType::I8 => cast_impl!(storage, i8),
|
||||
DType::U64 => cast_impl!(storage, u64),
|
||||
DType::U32 => cast_impl!(storage, u32),
|
||||
DType::U16 => cast_impl!(storage, u16),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
_ => panic!(
|
||||
_ => None,
|
||||
};
|
||||
let Some(bytes) = bytes else {
|
||||
panic!(
|
||||
"int_cast: unsupported conversion from {:?} to {:?}",
|
||||
tensor.dtype(),
|
||||
target_dtype
|
||||
),
|
||||
}
|
||||
)
|
||||
};
|
||||
FlexTensor::new(bytes, Layout::contiguous(shape), target_dtype)
|
||||
}
|
||||
|
||||
fn int_unfold(
|
||||
|
||||
Reference in New Issue
Block a user