feat:add triton op and cpp wrapper for copy_ and to_copy

ctests/CMakeLists.txt

@@ -98,3 +98,8 @@ add_executable(test_triton_rwkv_mm_sparsity test_triton_rwkv_mm_sparsity.cpp)
 target_link_libraries(test_triton_rwkv_mm_sparsity
   PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
 add_test(NAME test_triton_rwkv_mm_sparsity COMMAND test_triton_rwkv_mm_sparsity)
+
+add_executable(test_triton_copy test_triton_copy.cpp)
+target_link_libraries(test_triton_copy
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_copy COMMAND test_triton_copy)

ctests/test_triton_copy.cpp

@@ -0,0 +1,78 @@
+#include "flag_gems/operators.h"
+#include "gtest/gtest.h"
+#include "torch/torch.h"
+
+TEST(CopyTest, ContiguousTensorCopy) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor t = torch::randn({4, 5}, torch::TensorOptions().device(device).dtype(torch::kFloat32));
+
+  torch::Tensor out_gems = flag_gems::to_copy(t);
+  torch::Tensor out_ref = t.clone();
+
+  EXPECT_TRUE(torch::allclose(out_gems, out_ref));
+  EXPECT_EQ(out_gems.dtype(), t.dtype());
+}
+
+TEST(CopyTest, ContiguousTensorCopyWithDtype) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor t = torch::randn({3, 3}, torch::TensorOptions().device(device).dtype(torch::kFloat16));
+
+  torch::Tensor out_gems = flag_gems::to_copy(t, torch::kFloat32);
+  torch::Tensor out_ref = t.to(torch::kFloat32);
+
+  EXPECT_TRUE(torch::allclose(out_gems, out_ref));
+  EXPECT_EQ(out_gems.dtype(), torch::kFloat32);
+}
+
+TEST(CopyTest, NonContiguousTensorCopy) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor t = torch::randn({2, 3, 4}, torch::TensorOptions().device(device));
+  torch::Tensor t_transposed = t.transpose(0, 1);
+
+  torch::Tensor out_gems = flag_gems::to_copy(t_transposed);
+  torch::Tensor out_ref = t_transposed.clone();
+
+  EXPECT_TRUE(torch::allclose(out_gems, out_ref));
+}
+
+TEST(CopyTest, CopyInplaceContiguous) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor src = torch::randn({5, 5}, torch::TensorOptions().device(device));
+  torch::Tensor dst = torch::empty_like(src);
+
+  flag_gems::copy_(dst, src);
+
+  EXPECT_TRUE(torch::allclose(dst, src));
+}
+
+TEST(CopyTest, CopyInplaceNonContiguous) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor src = torch::randn({3, 4, 5}, torch::TensorOptions().device(device));
+  torch::Tensor dst = torch::empty({5, 4, 3}, torch::TensorOptions().device(device));
+  torch::Tensor src_transposed = src.transpose(0, 2);
+
+  flag_gems::copy_(dst, src_transposed);
+
+  EXPECT_TRUE(torch::allclose(dst, src_transposed));
+}
+
+TEST(CopyTest, CopyBroadcasting) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor src = torch::randn({1, 5}, torch::TensorOptions().device(device));
+  torch::Tensor dst = torch::empty({3, 5}, torch::TensorOptions().device(device));
+
+  flag_gems::copy_(dst, src);
+
+  torch::Tensor expected = src.expand_as(dst);
+  EXPECT_TRUE(torch::allclose(dst, expected));
+}
+
+TEST(CopyTest, EmptyTensor) {
+  const torch::Device device(torch::kCUDA, 0);
+  torch::Tensor src = torch::empty({0}, torch::TensorOptions().device(device));
+  torch::Tensor dst = torch::empty_like(src);
+
+  flag_gems::copy_(dst, src);
+
+  EXPECT_EQ(dst.numel(), 0);
+}

include/flag_gems/operators.h

@@ -226,4 +226,14 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> rwkv_ka_fusion(const at::Tensor &
                                                               int64_t H,
                                                               int64_t N);
 
+at::Tensor to_copy(const at::Tensor &self,
+                   c10::optional<at::ScalarType> dtype = c10::nullopt,
+                   c10::optional<at::Layout> layout = c10::nullopt,
+                   c10::optional<at::Device> device = c10::nullopt,
+                   c10::optional<bool> pin_memory = c10::nullopt,
+                   bool non_blocking = false,
+                   c10::optional<at::MemoryFormat> memory_format = c10::nullopt);
+
+at::Tensor &copy_(at::Tensor &dst, const at::Tensor &src, bool non_blocking = false);
+
 }  // namespace flag_gems

lib/CMakeLists.txt

@@ -27,7 +27,8 @@ add_library(operators
             reshape_and_cache_flash.cpp
             flash_attn_varlen_func.cpp
             rwkv_mm_sparsity.cpp
-            rwkv_ka_fusion.cpp)
+            rwkv_ka_fusion.cpp
+            copy.cpp)
 
 if (TRITON_GE_3P5)
   target_compile_definitions(operators PRIVATE TRITON_GE_3P5)

lib/copy.cpp

@@ -0,0 +1,232 @@
+#include <c10/core/DispatchKeySet.h>
+#include <vector>
+#include "c10/cuda/CUDAStream.h"
+#include "flag_gems/utils.h"
+#include "torch/torch.h"
+#include "triton_jit/triton_jit_function.h"
+
+namespace flag_gems {
+
+using namespace triton_jit;
+
+std::vector<int64_t> broadcasted_stride(const std::vector<int64_t>& shape,
+                                        const std::vector<int64_t>& stride,
+                                        const std::vector<int64_t>& target_shape) {
+  int ndim_diff = target_shape.size() - shape.size();
+  TORCH_CHECK(ndim_diff >= 0, "cannot broadcast to fewer dimensions");
+
+  std::vector<int64_t> full_shape(ndim_diff, 1);
+  full_shape.insert(full_shape.end(), shape.begin(), shape.end());
+
+  std::vector<int64_t> full_stride(ndim_diff, 0);
+  full_stride.insert(full_stride.end(), stride.begin(), stride.end());
+
+  std::vector<int64_t> out_stride(target_shape.size());
+
+  for (size_t i = 0; i < target_shape.size(); ++i) {
+    if (full_shape[i] == target_shape[i]) {
+      out_stride[i] = full_stride[i];
+    } else if (full_shape[i] == 1) {
+      out_stride[i] = 0;
+    } else {
+      TORCH_CHECK(false, "illegal broadcast at dim ", i);
+    }
+  }
+
+  return out_stride;
+}
+
+static bool _can_use_triton_copy(const at::Tensor& dst, const at::Tensor& src, bool non_blocking) {
+  if (!dst.is_cuda() || !src.is_cuda()) return false;
+  if (dst.device() != src.device()) return false;
+  if (non_blocking) return false;
+  return true;
+}
+
+static at::Tensor& redispatch_copy_fallback(at::Tensor& dst, const at::Tensor& src, bool non_blocking) {
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("aten::copy_", "")
+                       .typed<at::Tensor&(at::Tensor&, const at::Tensor&, bool)>();
+
+  constexpr c10::DispatchKeySet fallback_keyset =
+      c10::DispatchKeySet(c10::DispatchKey::CompositeExplicitAutograd);
+
+  return op.redispatch(fallback_keyset, dst, src, non_blocking);
+}
+
+static at::Tensor redispatch_to_copy_fallback(const at::Tensor& src,
+                                              c10::optional<at::ScalarType> dtype,
+                                              c10::optional<at::Layout> layout,
+                                              c10::optional<at::Device> device,
+                                              c10::optional<bool> pin_memory,
+                                              bool non_blocking,
+                                              c10::optional<at::MemoryFormat> memory_format) {
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("aten::_to_copy", "")
+                       .typed<at::Tensor(const at::Tensor&,
+                                         c10::optional<at::ScalarType>,
+                                         c10::optional<at::Layout>,
+                                         c10::optional<at::Device>,
+                                         c10::optional<bool>,
+                                         bool,
+                                         c10::optional<at::MemoryFormat>)>();
+
+  constexpr c10::DispatchKeySet fallback_keyset =
+      c10::DispatchKeySet(c10::DispatchKey::CompositeExplicitAutograd);
+
+  return op.redispatch(fallback_keyset, src, dtype, layout, device, pin_memory, non_blocking, memory_format);
+}
+
+at::Tensor to_copy(const at::Tensor& x,
+                   c10::optional<at::ScalarType> dtype = c10::nullopt,
+                   c10::optional<at::Layout> layout = c10::nullopt,
+                   c10::optional<at::Device> device = c10::nullopt,
+                   c10::optional<bool> pin_memory = c10::nullopt,
+                   bool non_blocking = false,
+                   c10::optional<at::MemoryFormat> memory_format = c10::nullopt) {
+  TORCH_WARN("[flag_gems][to_copy] gems::to_copy");
+  TORCH_CHECK(x.layout() == at::Layout::Strided, "Only strided tensors are supported");
+  TORCH_CHECK(!x.is_quantized(), "Quantized tensors are not supported");
+  if (layout.has_value()) {
+    TORCH_CHECK(layout.value() == x.layout(), "to_copy: layout conversion is not supported");
+  }
+  TORCH_CHECK(!pin_memory.has_value(), "to_copy: pin_memory is not supported");
+  TORCH_CHECK(!non_blocking, "to_copy: non_blocking copy is not supported");
+
+  auto target_dtype = dtype.has_value() ? dtype.value() : x.scalar_type();
+  auto target_device = device.has_value() ? device.value() : x.device();
+  auto target_memory_format = memory_format.has_value() ? memory_format.value() : at::MemoryFormat::Preserve;
+
+  at::Tensor out =
+      at::empty_like(x, x.options().dtype(target_dtype).device(target_device), target_memory_format);
+
+  // if (!_can_use_triton_copy(out, x, non_blocking)) {
+  //   return redispatch_to_copy_fallback(x, dtype, layout, device, pin_memory, non_blocking, memory_format);
+  // }
+
+  const int64_t numel = x.numel();
+  if (numel == 0) return out;
+
+  constexpr int BLOCK_SIZE = 1024;
+  const unsigned int grid_x = (numel + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+  c10::DeviceGuard guard(target_device);
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+
+  // at::Tensor x_linear = (x.scalar_type() != target_dtype) ? x.to(target_dtype) : x;
+  at::Tensor x_linear = x;
+  if (x.scalar_type() != target_dtype) {
+    return redispatch_to_copy_fallback(x, dtype, layout, device, pin_memory, non_blocking, memory_format);
+  }
+  if (x_linear.is_contiguous() && out.is_contiguous() && numel <= std::numeric_limits<int32_t>::max()) {
+    const TritonJITFunction& kernel_linear =
+        TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(),
+                                        "copy_kernel_linear");
+    kernel_linear(raw_stream, grid_x, 1, 1, 4, 0, x_linear, out, numel, BLOCK_SIZE);
+    return out;
+  }
+
+  std::vector<int64_t> task_shape(out.sizes().begin(), out.sizes().end());
+  int NDIMS = task_shape.size();
+
+  std::vector<int64_t> src_stride =
+      broadcasted_stride(std::vector<int64_t>(x_linear.sizes().begin(), x_linear.sizes().end()),
+                         std::vector<int64_t>(x_linear.strides().begin(), x_linear.strides().end()),
+                         task_shape);
+  std::vector<int64_t> dst_stride =
+      broadcasted_stride(std::vector<int64_t>(out.sizes().begin(), out.sizes().end()),
+                         std::vector<int64_t>(out.strides().begin(), out.strides().end()),
+                         task_shape);
+
+  const TritonJITFunction& kernel_nd =
+      TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(), "copy_kernel_nd");
+  kernel_nd(raw_stream,
+            grid_x,
+            1,
+            1,
+            4,
+            0,
+            x_linear,
+            out,
+            torch::tensor(task_shape, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
+            torch::tensor(src_stride, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
+            torch::tensor(dst_stride, torch::TensorOptions().dtype(torch::kInt64).device(out.device())),
+            numel,
+            NDIMS,
+            BLOCK_SIZE);
+
+  return out;
+}
+
+at::Tensor& copy_(at::Tensor& dst, const at::Tensor& src, bool non_blocking = false) {
+  TORCH_WARN("[flag_gems][copy_] gems::copy_");
+  if (!_can_use_triton_copy(dst, src, non_blocking)) {
+    return redispatch_copy_fallback(dst, src, non_blocking);
+  }
+  TORCH_CHECK(!dst._is_zerotensor(), "ZeroTensors are immutable");
+  if (src._is_zerotensor()) {
+    dst.zero_();
+    return dst;
+  }
+
+  if (dst.data_ptr() == src.data_ptr()) return dst;
+
+  TORCH_CHECK(src.sizes().size() <= dst.sizes().size(), "src cannot be broadcasted to dst");
+  for (size_t i = 0; i < src.dim(); ++i) {
+    TORCH_CHECK(src.size(i) == dst.size(dst.dim() - src.dim() + i) || src.size(i) == 1,
+                "src cannot be broadcasted to dst");
+  }
+
+  const int64_t numel = dst.numel();
+
+  constexpr int BLOCK_SIZE = 1024;
+  const unsigned int grid_x = (numel + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+  c10::DeviceGuard guard(dst.device());
+  c10::cuda::CUDAStream stream = c10::cuda::getCurrentCUDAStream();
+  CUstream raw_stream = static_cast<CUstream>(stream.stream());
+
+  bool no_broadcast = src.sizes().equals(dst.sizes());
+
+  if (dst.is_contiguous() && src.is_contiguous() && no_broadcast &&
+      numel <= std::numeric_limits<int32_t>::max()) {
+    const TritonJITFunction& kernel_linear =
+        TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(),
+                                        "copy_kernel_linear");
+    kernel_linear(raw_stream, grid_x, 1, 1, 4, 0, src, dst, numel, BLOCK_SIZE);
+    return dst;
+  }
+
+  std::vector<int64_t> task_shape(dst.sizes().begin(), dst.sizes().end());
+  int NDIMS = task_shape.size();
+
+  std::vector<int64_t> src_stride =
+      broadcasted_stride(std::vector<int64_t>(src.sizes().begin(), src.sizes().end()),
+                         std::vector<int64_t>(src.strides().begin(), src.strides().end()),
+                         task_shape);
+  std::vector<int64_t> dst_stride =
+      broadcasted_stride(std::vector<int64_t>(dst.sizes().begin(), dst.sizes().end()),
+                         std::vector<int64_t>(dst.strides().begin(), dst.strides().end()),
+                         task_shape);
+
+  const TritonJITFunction& kernel_nd =
+      TritonJITFunction::get_instance((utils::get_triton_src_path() / "copy.py").string(), "copy_kernel_nd");
+  kernel_nd(raw_stream,
+            grid_x,
+            1,
+            1,
+            4,
+            0,
+            src,
+            dst,
+            torch::tensor(task_shape, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
+            torch::tensor(src_stride, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
+            torch::tensor(dst_stride, torch::TensorOptions().dtype(torch::kInt64).device(dst.device())),
+            numel,
+            NDIMS,
+            BLOCK_SIZE);
+  return dst;
+}
+
+}  // namespace flag_gems

src/flag_gems/csrc/aten_patch.cpp

@@ -38,6 +38,8 @@ TORCH_LIBRARY_IMPL(aten, CUDA, m) {
   REGISTER_AND_LOG("zeros", zeros);
   REGISTER_AND_LOG("fill.Scalar", fill_scalar);
   REGISTER_AND_LOG("fill_.Scalar", fill_scalar_);
+  // REGISTER_AND_LOG("_to_copy", to_copy);
+  // REGISTER_AND_LOG("copy_", copy_);
 }
 
 }  // namespace flag_gems

src/flag_gems/csrc/cstub.cpp

@@ -36,6 +36,8 @@ PYBIND11_MODULE(c_operators, m) {
   m.def("remainder_.Tensor", &flag_gems::remainder_);
   m.def("rwkv_mm_sparsity", &flag_gems::rwkv_mm_sparsity);
   m.def("rwkv_ka_fusion", &flag_gems::rwkv_ka_fusion);
+  m.def("copy_", &flag_gems::copy_);
+  m.def("to_copy", &flag_gems::to_copy);
 }
 namespace flag_gems {
 TORCH_LIBRARY(flag_gems, m) {
@@ -123,6 +125,10 @@ TORCH_LIBRARY(flag_gems, m) {
 
   m.def("rwkv_mm_sparsity(Tensor k, Tensor v) -> Tensor");
   m.def("rwkv_ka_fusion(Tensor k, Tensor kk, Tensor a, Tensor ka, int H, int N) -> (Tensor, Tensor, Tensor)");
+  m.def("copy_(Tensor(a!) dst, Tensor src, bool non_blocking=False) -> Tensor(a!)");
+  m.def(
+      "to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? "
+      "pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor");
 }
 
 TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
@@ -195,5 +201,7 @@ TORCH_LIBRARY_IMPL(flag_gems, CUDA, m) {
   m.impl("flash_attn_varlen_func", TORCH_FN(flash_attn_varlen_func));
   m.impl("rwkv_mm_sparsity", TORCH_FN(rwkv_mm_sparsity));
   m.impl("rwkv_ka_fusion", TORCH_FN(rwkv_ka_fusion));
+  m.impl("to_copy", TORCH_FN(to_copy));
+  m.impl("copy_", TORCH_FN(copy_));
 }
 }  // namespace flag_gems