fix(infini): 为统一算子库修改

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

Cargo.toml

@@ -30,7 +30,7 @@ regex = "1.11"
 env_logger = "0.11"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "360b664", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "d22d2bc", default-features = false }
 search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "9b6289d" }
 search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "0e57976" }
 search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "5b9dbd9" }

models/llama/infini/src/infer.rs

@@ -3,7 +3,7 @@ use gguf::{ggml_quants::digit_layout::types, GGufModel};
 use llama::{ext::ggml_quants::f16, LlamaRequest, LlamaStorage, LlamaWorker, Tensor};
 use operators::{
     infini::Device,
-    infini_rt::{self, DeviceType::DEVICE_CPU},
+    infini_rt,
     random_sample::{KVPair, SampleArgs},
     TopoNode,
 };
@@ -59,7 +59,7 @@ fn test_infer() {
     let count = devices.len();
     println!("distribution: {devices:?}");
 
-    infini_rt::init(DEVICE_CPU);
+    infini_rt::init(infini_rt::DEVICE_CPU);
     let (seeds, senders) = WorkerSeed::new(devices.into_iter().map(|_| Device::cpu()).collect());
     thread::scope(|s| {
         let _workers = zip(lens, seeds)
@@ -90,8 +90,6 @@ fn test_infer() {
 
                     let sample = RandomSample::new(&node);
                     let indices = RandomSample::build_indices(model.meta.nvoc, &stream);
-                    let mut pair = KVPair::new(0, f16::ZERO);
-                    let mut pairs = Tensor::kv_pair_vec(1, |size| stream.malloc::<u8>(size));
 
                     for task in tasks {
                         let Task {
@@ -127,6 +125,11 @@ fn test_infer() {
                             )
                             .unwrap();
                         if id == 0 {
+                            // NOTICE 目前 random sample 完全是 CPU 上执行的，没必要再拷贝了
+                            let mut pair = KVPair::new(0, f16::ZERO);
+                            let mut pairs = Tensor::kv_pair_vec(1, |_| unsafe {
+                                from_raw_parts_mut(&mut pair as *mut _ as _, size_of_val(&pair))
+                            });
                             sample
                                 .launch(
                                     &mut pairs,
@@ -137,18 +140,7 @@ fn test_infer() {
                                     &stream,
                                 )
                                 .unwrap();
-
                             stream.synchronize();
-                            device.memcpy_d2h(
-                                unsafe {
-                                    from_raw_parts_mut(
-                                        &mut pair as *mut _ as *mut u8,
-                                        pairs.get().len(),
-                                    )
-                                },
-                                pairs.get(),
-                            );
-
                             next.send(pair.idx() as _).unwrap()
                         }
                     }

models/llama/infini/src/lib.rs

@@ -53,6 +53,7 @@ where
         let tensor = tensor.as_ref().map(|s| {
             let mut host = vec![0u8; s.len()];
             queue.get_device().memcpy_d2h(&mut host, s);
+            queue.synchronize();
             host
         });
         println!("{tensor}");

tensor/src/fmt.rs

@@ -56,6 +56,16 @@ impl DataFmt for u32 {
     }
 }
 
+impl DataFmt for u64 {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if *self == 0 {
+            write!(f, " ________")
+        } else {
+            write!(f, "{self:>6}")
+        }
+    }
+}
+
 impl<Physical: Deref<Target = [u8]>> fmt::Display for Tensor<Physical> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self.dt {
@@ -64,6 +74,7 @@ impl<Physical: Deref<Target = [u8]>> fmt::Display for Tensor<Physical> {
             primitive::F32 => self.map_slice().write_tensor::<f32>(&mut vec![], f),
             primitive::F64 => self.map_slice().write_tensor::<f64>(&mut vec![], f),
             primitive::U32 => self.map_slice().write_tensor::<u32>(&mut vec![], f),
+            primitive::U64 => self.map_slice().write_tensor::<u64>(&mut vec![], f),
             _ => todo!(),
         }
     }