From 7950ff847541de8e5d19a90e3b33b34656ef8437 Mon Sep 17 00:00:00 2001
From: Kevin Oberlies <favilo@gmail.com>
Date: Sat, 2 Dec 2023 18:36:40 -0800
Subject: [PATCH] no-std

---
 dfdx-core/src/tensor/webgpu/device.rs | 28 ++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/dfdx-core/src/tensor/webgpu/device.rs b/dfdx-core/src/tensor/webgpu/device.rs
index 4eb62c97..086eb4e4 100644
--- a/dfdx-core/src/tensor/webgpu/device.rs
+++ b/dfdx-core/src/tensor/webgpu/device.rs
@@ -13,18 +13,13 @@ use crate::{
 };
 
 #[cfg(feature = "no-std")]
-use spin::Mutex;
+use spin::{Mutex, RwLock};
 
 use core::any::TypeId;
 #[cfg(not(feature = "no-std"))]
-use std::sync::Mutex;
+use std::sync::{Mutex, RwLock};
 
-use std::{
-    collections::HashMap,
-    marker::PhantomData,
-    sync::{Arc, RwLock},
-    vec::Vec,
-};
+use std::{collections::HashMap, marker::PhantomData, sync::Arc, vec::Vec};
 
 use super::allocate::round_to_buffer_alignment;
 
@@ -209,22 +204,37 @@ impl Webgpu {
         Ok(data)
     }
 
+    #[cfg(not(feature = "no-std"))]
     pub(crate) fn shader_module_loaded(&self, name: TypeId) -> bool {
         self.cs_cache.read().unwrap().contains_key(&name)
     }
 
+    #[cfg(feature = "no-std")]
+    pub(crate) fn shader_module_loaded(&self, name: TypeId) -> bool {
+        self.cs_cache.read().contains_key(&name)
+    }
+
     pub(crate) fn load_shader_module(&self, name: TypeId, source: &str) {
         let module = Arc::new(self.dev.create_shader_module(ShaderModuleDescriptor {
             label: None,
             source: wgpu::ShaderSource::Wgsl(source.into()),
         }));
+        #[cfg(not(feature = "no-std"))]
         self.cs_cache.write().unwrap().insert(name, module);
+        #[cfg(feature = "no-std")]
+        self.cs_cache.write().insert(name, module);
     }
 
+    #[cfg(not(feature = "no-std"))]
     pub(crate) fn get_shader_module(&self, name: TypeId) -> Option<Arc<ShaderModule>> {
         self.cs_cache.read().unwrap().get(&name).cloned()
     }
 
+    #[cfg(feature = "no-std")]
+    pub(crate) fn get_shader_module(&self, name: TypeId) -> Option<Arc<ShaderModule>> {
+        self.cs_cache.read().get(&name).cloned()
+    }
+
     // #[allow(unused)]
     // pub(crate) unsafe fn get_workspace<E>(&self, len: usize) -> Result<MutexGuard<Buffer>, Error> {
     //     let num_bytes_required = len * std::mem::size_of::<E>();
@@ -369,7 +379,7 @@ impl<E: Unit> Storage<E> for Webgpu {
     type Vec = CachableBuffer<E>;
 
     fn try_alloc_len(&self, len: usize) -> Result<Self::Vec, Error> {
-        let data = unsafe { self.alloc_empty::<E>(len) }?;
+        let data = self.alloc_empty::<E>(len)?;
         Ok(CachableBuffer {
             dev: self.dev.clone(),
             queue: self.queue.clone(),