feat(flow-ml): real hanzo-ml binding — GpuResident seam made concrete

MlDevice wraps hanzo_ml::Device (CPU/CUDA/Metal/…); MlTensor wraps hanzo_ml::Tensor — already an Arc-backed device-resident handle, so it IS the DeviceTensor abstraction concretely (zero-copy step-to-step). The same steps run on CPU or GPU; only the Device ctor differs (MlDevice::cpu vs ::cuda). flow-core stays std-only — hanzo-ml lives behind this crate so hanzo-desktop embeds flow-core alone (tiny) or flow-ml (with compute) as needed. Verified on CPU (no GPU): a real affine op runs through GpuResidentBackend, 1.0→2.0, result resident on the injected device. 2/2, clippy clean. cuda/metal features forward to hanzo-ml for on-GPU residence.
2026-10-17 01:29:18 +00:00 · 2026-07-16 19:21:15 -07:00
parent 9f740023f9
commit a5397cc135
3 changed files with 204 additions and 0 deletions
@@ -0,0 +1,21 @@
+[package]
+name = "flow-ml"
+version = "0.1.0"
+edition = "2021"
+rust-version = "1.88"
+description = "The hanzo-ml binding for flow-core: MlDevice/MlTensor make the GpuResident seam concrete. Same steps run on CPU or GPU — only the Device ctor differs. CPU by default; enable `cuda`/`metal` for on-GPU residence."
+license = "MIT"
+
+[lib]
+name = "flow_ml"
+
+# flow-core stays std-only; the heavy compute dep lives here, behind this crate, so
+# hanzo-desktop can embed flow-core alone (tiny) or flow-ml (with hanzo-ml) as needed.
+[dependencies]
+flow-core = { path = "../core-rs" }
+hanzo-ml = { path = "../../ml/hanzo-ml" }
+
+[features]
+# CPU by default (default = []). A GPU-resident build just forwards to hanzo-ml.
+cuda = ["hanzo-ml/cuda"]
+metal = ["hanzo-ml/metal"]
@@ -0,0 +1,92 @@
+//! `flow-ml` — the real `hanzo-ml` binding for flow-core's `GpuResident` seam.
+//!
+//! [`MlDevice`] wraps a `hanzo_ml::Device` (CPU / CUDA / Metal / Vulkan / ROCm) as a
+//! [`flow_core::gpu::Device`]. [`MlTensor`] wraps a `hanzo_ml::Tensor` as a
+//! [`flow_core::Value`] — and a `hanzo_ml::Tensor` is *already* an `Arc`-backed,
+//! device-resident handle, so this is the [`flow_core::gpu::DeviceTensor`] abstraction
+//! made concrete: cloning it between steps is a refcount bump, never a host copy.
+//!
+//! The payoff of the seam: the *same* step code runs on CPU or GPU — only the `Device`
+//! constructor differs ([`MlDevice::cpu`] vs [`MlDevice::cuda`]). Tests run on CPU (no
+//! GPU, no CUDA build); production runs the identical graph on a CUDA/Metal device.
+
+use std::any::Any;
+
+use flow_core::gpu::Device as FlowDevice;
+use flow_core::Value;
+use hanzo_ml::{Device as MlDev, Tensor};
+
+/// A flow-core [`Device`](flow_core::gpu::Device) backed by a real `hanzo_ml::Device`.
+/// GPU steps recover it via `current_device().and_then(|d| d.as_any().downcast_ref())`
+/// and reach the underlying `hanzo_ml::Device` with [`MlDevice::device`] to run kernels.
+pub struct MlDevice {
+    inner: MlDev,
+    name: String,
+}
+
+impl MlDevice {
+    pub fn new(inner: MlDev) -> Self {
+        let name = match inner.location() {
+            hanzo_ml::DeviceLocation::Cpu => "cpu".to_string(),
+            hanzo_ml::DeviceLocation::Cuda { gpu_id } => format!("cuda:{gpu_id}"),
+            hanzo_ml::DeviceLocation::Metal { gpu_id } => format!("metal:{gpu_id}"),
+        };
+        Self { inner, name }
+    }
+    /// The CPU device — always available, no GPU required (what tests run on).
+    pub fn cpu() -> Self {
+        Self::new(MlDev::Cpu)
+    }
+    /// A CUDA device by ordinal. Requires the `cuda` feature; the identical graph that
+    /// ran on [`MlDevice::cpu`] runs here on-GPU.
+    pub fn cuda(ordinal: usize) -> hanzo_ml::Result<Self> {
+        Ok(Self::new(MlDev::new_cuda(ordinal)?))
+    }
+    /// A CUDA device if one is available, else CPU — the desktop/BYO-GPU fallback path.
+    pub fn cuda_if_available(ordinal: usize) -> hanzo_ml::Result<Self> {
+        Ok(Self::new(MlDev::cuda_if_available(ordinal)?))
+    }
+    /// The underlying `hanzo_ml::Device`, for constructing/executing tensors on it.
+    pub fn device(&self) -> &MlDev {
+        &self.inner
+    }
+}
+
+impl FlowDevice for MlDevice {
+    fn name(&self) -> &str {
+        &self.name
+    }
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
+/// A flow-core [`Value`] wrapping a device-resident `hanzo_ml::Tensor`. Its `port_type`
+/// is `"tensor"` — the type GPU ports carry. Cloning is a refcount bump (the tensor is
+/// `Arc`-backed), so it flows step-to-step zero-copy and never round-trips to host.
+#[derive(Clone)]
+pub struct MlTensor(pub Tensor);
+
+impl MlTensor {
+    pub fn into_inner(self) -> Tensor {
+        self.0
+    }
+    pub fn tensor(&self) -> &Tensor {
+        &self.0
+    }
+}
+
+impl From<Tensor> for MlTensor {
+    fn from(t: Tensor) -> Self {
+        MlTensor(t)
+    }
+}
+
+impl Value for MlTensor {
+    fn port_type(&self) -> &str {
+        "tensor"
+    }
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
@@ -0,0 +1,91 @@
+//! Proves the real hanzo-ml binding end-to-end on CPU: real `hanzo_ml::Tensor`s flow
+//! through flow-core's `GpuResidentBackend` as values, a real on-device op (`affine`)
+//! executes, and the numeric result is correct. No GPU, no CUDA build — the identical
+//! graph runs on a CUDA/Metal device by constructing `MlDevice::cuda(0)` instead.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use flow_core::gpu::{current_device, GpuResidentBackend};
+use flow_core::*;
+use flow_ml::{MlDevice, MlTensor};
+use hanzo_ml::{DType, Device as MlDev, Tensor};
+
+fn tensor_of(v: &ValueRef) -> Tensor {
+    v.as_any().downcast_ref::<MlTensor>().expect("MlTensor").0.clone()
+}
+
+/// Recover the real hanzo_ml::Device the step is running on — what a GPU step does to
+/// build/execute tensors on the injected device.
+fn ml_device() -> MlDev {
+    let d = current_device().expect("device injected in GpuResident mode");
+    d.as_any().downcast_ref::<MlDevice>().expect("MlDevice").device().clone()
+}
+
+// Source: allocates a [2,3] ones tensor on the injected device (host→device once).
+struct Ones {
+    id: String,
+}
+impl Step for Ones {
+    fn id(&self) -> &str {
+        &self.id
+    }
+    fn ins(&self) -> Vec<Port> {
+        vec![]
+    }
+    fn outs(&self) -> Vec<Port> {
+        vec![Port::new("t", "tensor")]
+    }
+    fn run(&self, _in: &Inputs) -> Result<Outputs, FlowError> {
+        let dev = ml_device();
+        let t = Tensor::ones((2, 3), DType::F32, &dev)
+            .map_err(|e| FlowError::Step { id: self.id.clone(), source: e.to_string() })?;
+        Ok(HashMap::from([("t".to_string(), Arc::new(MlTensor(t)) as ValueRef)]))
+    }
+}
+
+// Transform: a real on-device op — affine(2, 0) doubles every element, staying resident.
+struct Double {
+    id: String,
+}
+impl Step for Double {
+    fn id(&self) -> &str {
+        &self.id
+    }
+    fn ins(&self) -> Vec<Port> {
+        vec![Port::new("t", "tensor")]
+    }
+    fn outs(&self) -> Vec<Port> {
+        vec![Port::new("t", "tensor")]
+    }
+    fn run(&self, inp: &Inputs) -> Result<Outputs, FlowError> {
+        let t = tensor_of(inp.get("t").expect("t"));
+        let out = t
+            .affine(2.0, 0.0)
+            .map_err(|e| FlowError::Step { id: self.id.clone(), source: e.to_string() })?;
+        Ok(HashMap::from([("t".to_string(), Arc::new(MlTensor(out)) as ValueRef)]))
+    }
+}
+
+#[test]
+fn ml_pipeline_runs_a_real_op_on_cpu() {
+    let mut g = Graph::new();
+    g.add(Arc::new(Ones { id: "ones".into() })).unwrap();
+    g.add(Arc::new(Double { id: "double".into() })).unwrap();
+    g.wire(PortRef::new("ones", "t"), PortRef::new("double", "t")).unwrap();
+
+    let sched = Scheduler::new().with_backend(Box::new(GpuResidentBackend::new(Arc::new(MlDevice::cpu()))));
+    let out = sched.run(&g, Mode::GpuResident, &HashMap::new()).unwrap();
+
+    let t = tensor_of(out.get("double.t").expect("leaf"));
+    assert_eq!(t.dims(), &[2, 3]);
+    // The real op ran: every 1.0 became 2.0.
+    assert_eq!(t.to_vec2::<f32>().unwrap(), vec![vec![2.0, 2.0, 2.0], vec![2.0, 2.0, 2.0]]);
+    // And it ran on the CPU device we injected (swap MlDevice::cuda(0) → runs on GPU).
+    assert!(matches!(t.device().location(), hanzo_ml::DeviceLocation::Cpu));
+}
+
+#[test]
+fn ml_device_names_reflect_location() {
+    assert_eq!(<MlDevice as flow_core::gpu::Device>::name(&MlDevice::cpu()), "cpu");
+}