Replace MetalConst -> Const

sonos · Dec 6, 2024 · 024cffc · 024cffc
1 parent 3df1f74
commit 024cffc
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 49 deletions.
diff --git a/metal/src/fact.rs b/metal/src/fact.rs
@@ -27,7 +27,7 @@ impl MetalFact {
     }
 
     pub fn from_cpu(fact: TypedFact) -> TractResult<Self> {
-        Self::new(MetalOrigin::FromGpu, fact)
+        Self::new(MetalOrigin::FromCpu, fact)
     }
 
     pub fn is_from_gpu(&self) -> bool {

diff --git a/metal/src/ops/konst.rs b/metal/src/ops/konst.rs
diff --git a/metal/src/ops/mod.rs b/metal/src/ops/mod.rs
@@ -6,7 +6,6 @@ pub mod change_axes;
 pub mod concat;
 pub mod element_wise;
 pub mod gemm;
-pub mod konst;
 pub mod new_gelu;
 pub mod reduce;
 pub mod rms_norm;
@@ -24,7 +23,6 @@ pub use change_axes::{MetalAxisOp, MetalIntoShape};
 pub use concat::MetalConcat;
 pub use element_wise::MetalElementWiseOp;
 pub use gemm::MetalGemm;
-pub use konst::MetalConst;
 pub use new_gelu::MetalNewGelu;
 pub use reduce::MetalReduce;
 pub use rms_norm::MetalRmsNorm;

diff --git a/metal/src/transform.rs b/metal/src/transform.rs
@@ -184,7 +184,7 @@ impl Translate<TypedFact, Box<dyn TypedOp>, TypedFact, Box<dyn TypedOp>> for Met
         } else if let Some(op) = node.op_as::<MultiBroadcastTo>() {
             Some(Box::new(ops::MetalMultiBroadcastTo::new(op.shape.clone())))
         } else if let Some(op) = node.op_as::<Const>() {
-            ops::MetalConst::new(op.0.clone())?.map(|o| -> Box<dyn TypedOp> { Box::new(o) })
+            convert_const(op)?.map(|o| -> Box<dyn TypedOp> { Box::new(o) })
         } else if let Some(op) = node.op_as::<Cast>() {
             check_in_dts_are_supported(source, node.id, ops::MetalCast::is_supported_dt)?
                 .then(|| ops::MetalCast::new(op.to))
@@ -355,6 +355,16 @@ pub fn bin_ops_to_metal(
         .transpose()
 }
 
+fn convert_const(op: &Const) -> TractResult<Option<Const>> {
+    if !MetalTensor::is_supported_dt(op.0.datum_type()) {
+        return Ok(None);
+    }
+    let metal_fact = MetalFact::from_cpu(Arc::clone(&op.0).into())?;
+    let metal_const = op.0.clone().into_metal()?.into_opaque_tensor().into_arc_tensor();
+    Ok(Some(Const::new_with_opaque_fact(metal_const, Box::new(metal_fact))))
+}
+
+
 fn convert_element_wise_ops_to_metal(op: &ElementWiseOp) -> Option<ops::MetalElementWiseOp> {
     map_element_wise_ops!([
         (tract_core::ops::math::Abs, Abs),