From 634a68932717847ef0253267e3b22217a2b827b6 Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Wed, 7 Aug 2024 09:55:03 -0500 Subject: [PATCH] fix: add other allocation devices. closes #253 (cherry picked from commit 08aaa0ffabf597eb97514180460de65bd1bf9938) --- src/memory.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/memory.rs b/src/memory.rs index 74eb7703..b3ea48ae 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -145,7 +145,10 @@ pub enum AllocationDevice { HIP, HIPPinned, OpenVINOCPU, - OpenVINOGPU + OpenVINOGPU, + // these aren't defined in allocator.h and are instead scattered all over. fun! + DirectMLCPU, + TVM } impl AllocationDevice { @@ -157,17 +160,19 @@ impl AllocationDevice { Self::CANN => "Cann", Self::CANNPinned => "CannPinned", Self::DirectML => "Dml", + Self::DirectMLCPU => "DML CPU", // yes, caps & space Self::HIP => "Hip", Self::HIPPinned => "HipPinned", Self::OpenVINOCPU => "OpenVINO_CPU", - Self::OpenVINOGPU => "OpenVINO_GPU" + Self::OpenVINOGPU => "OpenVINO_GPU", + Self::TVM => "TVM" } } /// Returns `true` if this memory is accessible by the CPU; meaning that, if a value were allocated on this device, /// it could be extracted to an `ndarray` or slice. pub fn is_cpu_accessible(&self) -> bool { - matches!(self, Self::CPU | Self::CUDAPinned | Self::CANNPinned | Self::HIPPinned | Self::OpenVINOCPU) + matches!(self, Self::CPU | Self::CUDAPinned | Self::CANNPinned | Self::HIPPinned | Self::OpenVINOCPU | Self::DirectMLCPU | Self::TVM) } } @@ -182,10 +187,12 @@ impl TryFrom for AllocationDevice { "Cann" => Ok(AllocationDevice::CANN), "CannPinned" => Ok(AllocationDevice::CANNPinned), "Dml" => Ok(AllocationDevice::DirectML), + "DML CPU" => Ok(AllocationDevice::DirectMLCPU), "Hip" => Ok(AllocationDevice::HIP), "HipPinned" => Ok(AllocationDevice::HIPPinned), "OpenVINO_CPU" => Ok(AllocationDevice::OpenVINOCPU), "OpenVINO_GPU" => Ok(AllocationDevice::OpenVINOGPU), + "TVM" => Ok(AllocationDevice::TVM), _ => Err(value) } }