Update CopyKernel.cpp

Author: yucai-intel

src/ATen/native/xpu/sycl/CopyKernel.cpp

@@ -5,9 +5,10 @@
 #include <ATen/core/Tensor.h>
 #include <ATen/native/TensorIterator.h>
 
-#include <ATen/native/xpu/sycl/CopyKernel.h>
 #include <ATen/native/xpu/sycl/Loops.h>
 
+#include <ATen/native/xpu/sycl/CopyKernel.h>
+
 namespace at::native::xpu {
 
 template <typename scalar_t>
@@ -101,6 +102,21 @@ void float8_copy_kernel_xpu(TensorIteratorBase& iter) {
         gpu_kernel(iter, CopyScalarFunc<Float8_e5m2fnuz>());
         break;
     }
+  } else if (dtype == kFloat8_e8m0fnu) {
+    switch (other_dtype) {
+      case kFloat:
+        gpu_kernel_nocast(iter, CastScalarFunc<float, Float8_e8m0fnu>());
+        break;
+      case kHalf:
+        gpu_kernel_nocast(iter, CastScalarFunc<Half, Float8_e8m0fnu>());
+        break;
+      case kBFloat16:
+        gpu_kernel_nocast(iter, CastScalarFunc<BFloat16, Float8_e8m0fnu>());
+        break;
+      default:
+        gpu_kernel(iter, CopyScalarFunc<Float8_e8m0fnu>());
+        break;
+    }
   } else {
     TORCH_CHECK(
         false,
@@ -109,6 +125,16 @@ void float8_copy_kernel_xpu(TensorIteratorBase& iter) {
   }
 }
 
+void float4_copy_kernel_xpu(TensorIteratorBase& iter) {
+  ScalarType src_dtype = iter.dtype(1);
+
+  if (src_dtype == kFloat4_e2m1fn_x2) {
+    gpu_kernel_nocast(iter, CopyScalarFunc<Float4_e2m1fn_x2>());
+  } else {
+    TORCH_CHECK(false, "Copy from ", src_dtype, " to Float4_e2m1fn_x2 has not been supported.");
+  }
+}
+
 void copy_kernel(TensorIteratorBase& iter) {
   ScalarType dtype = iter.common_dtype();
   if (isQIntType(dtype)) {
@@ -117,6 +143,8 @@ void copy_kernel(TensorIteratorBase& iter) {
     });
   } else if (isFloat8Type(iter.dtype(0))) {
     float8_copy_kernel_xpu(iter);
+  } else if (iter.dtype(0) == kFloat4_e2m1fn_x2) {
+    float4_copy_kernel_xpu(iter);
   } else {
     AT_DISPATCH_V2(
         dtype,
@@ -127,11 +155,8 @@ void copy_kernel(TensorIteratorBase& iter) {
         kBool,
         kBFloat16,
         kComplexHalf,
-        AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES),
-        kFloat8_e4m3fn,
-        kFloat8_e5m2,
-        kFloat8_e4m3fnuz,
-        kFloat8_e5m2fnuz);
+        AT_EXPAND(AT_FLOAT8_TYPES),
+        AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
   }
 }