Fix alpha handling of GPU reformatter

Author: tongyuantongyu

reformat/reformat.cu

@@ -127,13 +127,8 @@ void import_pixel_cuda(md_view<F, 3> dst,
     import_opaque_kernel<<<dimGrid, dimBlock, 0, stream>>>(dst, src, F(a), F(b));
   }
   else {
-    if (c == 4) {
-      import_alpha_kernel<<<dimGrid, dimBlock, 0, stream>>>(dst, dst_alpha, src, F(a), F(b));
-    }
-    else {
-      import_opaque_kernel<<<dimGrid, dimBlock, 0, stream>>>(dst, src, F(a), F(b));
-      cudaMemsetAsync(dst_alpha.data, 0, dst_alpha.size() * sizeof(F), stream);
-    }
+    assert(src.shape[2] == 4);
+    import_alpha_kernel<<<dimGrid, dimBlock, 0, stream>>>(dst, dst_alpha, src, F(a), F(b));
   }
 }
 

reformat/reformat.h

@@ -218,6 +218,30 @@ std::string pixel_importer_cpu::import_alpha(md_view<float, 3> dst,
   return "";
 }
 
+
+template<typename T>
+std::string pixel_exporter_cpu::fetch_alpha(md_view<const float, 3> src, cudaStream_t stream) {
+  auto [c, h, w] = src.shape;
+
+  if (h * w > max_size) {
+    return "dimension too big";
+  }
+
+  auto err = cudaMemcpyAsync(buffer_alpha.get(), src.data, h * w * 4, cudaMemcpyDeviceToHost, stream);
+  if (err != cudaSuccess) {
+    return std::string("CUDA error: ") + cudaGetErrorName(err);
+  }
+
+  err = cudaStreamSynchronize(stream);
+  if (err != cudaSuccess) {
+    return std::string("CUDA error: ") + cudaGetErrorName(err);
+  }
+
+  current_buffer_shape = src.shape;
+  alpha_filled = true;
+  return "";
+}
+
 template<std::unsigned_integral U>
 std::string pixel_exporter_cpu::fetch_color(md_view<const float, 3> src,
                                             md_uview<U, 3> dst,
@@ -320,29 +344,6 @@ std::string pixel_exporter_cpu::fetch_color(md_view<const float, 3> src,
   return "";
 }
 
-template<typename T>
-std::string pixel_exporter_cpu::fetch_alpha(md_view<const float, 3> src, cudaStream_t stream) {
-  auto [c, h, w] = src.shape;
-
-  if (h * w > max_size) {
-    return "dimension too big";
-  }
-
-  auto err = cudaMemcpyAsync(buffer_alpha.get(), src.data, h * w * 4, cudaMemcpyDeviceToHost, stream);
-  if (err != cudaSuccess) {
-    return std::string("CUDA error: ") + cudaGetErrorName(err);
-  }
-
-  err = cudaStreamSynchronize(stream);
-  if (err != cudaSuccess) {
-    return std::string("CUDA error: ") + cudaGetErrorName(err);
-  }
-
-  current_buffer_shape = src.shape;
-  alpha_filled = true;
-  return "";
-}
-
 // -----------------------------------------------------------------------------
 // GPU part
 
@@ -523,7 +524,7 @@ std::string pixel_exporter_gpu<F, eSize>::fetch_alpha(md_view<const F, 3> src, c
     return "dimension too big";
   }
 
-  auto err = cudaMemcpyAsync(gpu_buffer_alpha, src.data, h * w, cudaMemcpyDeviceToDevice, stream);
+  auto err = cudaMemcpyAsync(gpu_buffer_alpha, src.data, h * w * sizeof(F), cudaMemcpyDeviceToDevice, stream);
   if (err != cudaSuccess) {
     return std::string("CUDA error: ") + cudaGetErrorName(err);
   }