Fix discarding image data when transitioning StorageImage (#2150)

* Fix discarding image data when transitioning StorageImage (#2149)

* Used Relaxed ordering for `StorageImage::layout_initialized` atomic

vulkano/src/image/storage.rs

@@ -32,7 +32,10 @@ use smallvec::SmallVec;
 use std::{
     fs::File,
     hash::{Hash, Hasher},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
 };
 
 /// General-purpose image in device memory. Can be used for any usage, but will be slower than a
@@ -40,6 +43,10 @@ use std::{
 #[derive(Debug)]
 pub struct StorageImage {
     inner: Arc<Image>,
+
+    // If true, then the image is in the layout `General`. If false, then it
+    // is still `Undefined`.
+    layout_initialized: AtomicBool,
 }
 
 impl StorageImage {
@@ -127,7 +134,10 @@ impl StorageImage {
                         .map_err(|(err, _, _)| err)?
                 });
 
-                Ok(Arc::new(StorageImage { inner }))
+                Ok(Arc::new(StorageImage {
+                    inner,
+                    layout_initialized: AtomicBool::new(false),
+                }))
             }
             Err(err) => Err(err.into()),
         }
@@ -204,7 +214,10 @@ impl StorageImage {
                         .map_err(|(err, _, _)| err)?
                 });
 
-                Ok(Arc::new(StorageImage { inner }))
+                Ok(Arc::new(StorageImage {
+                    inner,
+                    layout_initialized: AtomicBool::new(false),
+                }))
             }
             Err(err) => Err(err.into()),
         }
@@ -301,6 +314,16 @@ unsafe impl ImageAccess for StorageImage {
         ImageLayout::General
     }
 
+    #[inline]
+    unsafe fn layout_initialized(&self) {
+        self.layout_initialized.store(true, Ordering::Relaxed);
+    }
+
+    #[inline]
+    fn is_layout_initialized(&self) -> bool {
+        self.layout_initialized.load(Ordering::Relaxed)
+    }
+
     #[inline]
     fn descriptor_layouts(&self) -> Option<ImageDescriptorLayouts> {
         Some(ImageDescriptorLayouts {