feat(cambricon): add copy provider

src/native/cambricon/cnnl_utils.h

@@ -3,14 +3,19 @@
 
 #include <cassert>
 #include <cstdint>
-#include <limits>
 #include <memory>
-#include <type_traits>
 #include <vector>
 
 #include "native/cambricon/common.h"
 #include "tensor.h"
 
+#define INFINI_OPS_CNNL_CHECK(call)                                      \
+  do {                                                                   \
+    const auto cnnl_status = (call);                                     \
+    assert(cnnl_status == CNNL_STATUS_SUCCESS && "`" #call "` failed."); \
+    (void)cnnl_status;                                                   \
+  } while (false)
+
 namespace infini::ops::cnnl_utils {
 
 struct HandleDeleter {
@@ -28,8 +33,7 @@ using Handle =
 
 inline Handle CreateHandle() {
   cnnlHandle_t handle{nullptr};
-  [[maybe_unused]] const auto status = cnnlCreate(&handle);
-  assert(status == CNNL_STATUS_SUCCESS && "`cnnlCreate` failed.");
+  INFINI_OPS_CNNL_CHECK(cnnlCreate(&handle));
 
   return Handle{handle};
 }
@@ -50,75 +54,41 @@ using TensorDescriptor =
 
 inline TensorDescriptor CreateTensorDescriptor() {
   cnnlTensorDescriptor_t desc{nullptr};
-  [[maybe_unused]] const auto status = cnnlCreateTensorDescriptor(&desc);
-  assert(status == CNNL_STATUS_SUCCESS &&
-         "`cnnlCreateTensorDescriptor` failed.");
+  INFINI_OPS_CNNL_CHECK(cnnlCreateTensorDescriptor(&desc));
 
   return TensorDescriptor{desc};
 }
 
-namespace detail {
-
-template <typename Integer>
-int CheckedInt(Integer value) {
-  static_assert(std::is_integral_v<Integer>);
-
-  [[maybe_unused]] bool out_of_range{false};
-  if constexpr (std::is_signed_v<Integer>) {
-    const auto wide = static_cast<std::intmax_t>(value);
-    out_of_range = wide < std::numeric_limits<int>::min() ||
-                   wide > std::numeric_limits<int>::max();
-  } else {
-    const auto wide = static_cast<std::uintmax_t>(value);
-    out_of_range =
-        wide > static_cast<std::uintmax_t>(std::numeric_limits<int>::max());
-  }
-
-  assert(!out_of_range &&
-         "`CNNL tensor descriptor` value does not fit in `int`.");
-
-  return static_cast<int>(value);
-}
-
-template <typename Values>
-std::vector<int> CheckedIntVector(const Values& values) {
-  std::vector<int> result;
-  result.reserve(values.size());
-  for (const auto value : values) {
-    result.push_back(CheckedInt(value));
-  }
-  return result;
-}
-
-}  // namespace detail
-
 inline void SetTensorDescriptor(cnnlTensorDescriptor_t desc, DataType dtype,
                                 const Tensor::Shape& shape,
                                 const Tensor::Strides& strides) {
-  assert(!shape.empty() && shape.size() == strides.size() &&
-         "`CNNL tensor descriptor` requires matching non-empty shape and "
-         "strides.");
+  assert(shape.size() == strides.size() &&
+         "`CNNL tensor descriptor` requires matching shape and strides.");
 
   const auto cnnl_dtype = GetDataType(dtype);
   assert(cnnl_dtype != CNNL_DTYPE_INVALID &&
          "`CNNL tensor descriptor` does not support this data type.");
 
-  const auto cnnl_shape = detail::CheckedIntVector(shape);
-  const auto cnnl_strides = detail::CheckedIntVector(strides);
-  const auto ndim = detail::CheckedInt(shape.size());
-
-  [[maybe_unused]] const auto status =
-      cnnlSetTensorDescriptorEx(desc, CNNL_LAYOUT_ARRAY, cnnl_dtype, ndim,
-                                cnnl_shape.data(), cnnl_strides.data());
-  assert(status == CNNL_STATUS_SUCCESS &&
-         "`cnnlSetTensorDescriptorEx` failed.");
+  const auto cnnl_shape =
+      shape.empty() ? std::vector<std::int64_t>{1}
+                    : std::vector<std::int64_t>(shape.begin(), shape.end());
+  const auto cnnl_strides =
+      strides.empty()
+          ? std::vector<std::int64_t>{1}
+          : std::vector<std::int64_t>(strides.begin(), strides.end());
+  const auto ndim = static_cast<int>(cnnl_shape.size());
+
+  INFINI_OPS_CNNL_CHECK(
+      cnnlSetTensorDescriptorEx_v2(desc, CNNL_LAYOUT_ARRAY, cnnl_dtype, ndim,
+                                   cnnl_shape.data(), cnnl_strides.data()));
 }
 
 inline TensorDescriptor MakeTensorDescriptor(DataType dtype,
                                              const Tensor::Shape& shape,
                                              const Tensor::Strides& strides) {
   auto desc = CreateTensorDescriptor();
   SetTensorDescriptor(desc.get(), dtype, shape, strides);
+
   return desc;
 }
 

src/native/cambricon/ops/copy/cnnl.h

@@ -0,0 +1,80 @@
+#ifndef INFINI_OPS_CAMBRICON_COPY_CNNL_H_
+#define INFINI_OPS_CAMBRICON_COPY_CNNL_H_
+
+#include <algorithm>
+#include <cassert>
+
+#include "base/copy.h"
+#include "native/cambricon/cnnl_utils.h"
+#include "native/cambricon/cnrt_utils.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Copy, Device::Type::kCambricon> : public Copy {
+ public:
+  Operator(const Tensor src, const bool non_blocking, Tensor out)
+      : Copy{src, non_blocking, out} {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    assert(std::all_of(input_strides_.begin(), input_strides_.end(),
+                       [](auto stride) { return stride >= 0; }) &&
+           "`CambriconCopy` does not support negative input strides.");
+    assert(std::all_of(out_strides_.begin(), out_strides_.end(),
+                       [](auto stride) { return stride >= 0; }) &&
+           "`CambriconCopy` does not support negative output strides.");
+
+    cnnl_handle_ = cnnl_utils::CreateHandle();
+    input_desc_ = cnnl_utils::MakeTensorDescriptor(input_type_, input_shape_,
+                                                   input_strides_);
+    out_desc_ =
+        cnnl_utils::MakeTensorDescriptor(out_type_, out_shape_, out_strides_);
+
+    INFINI_OPS_CNNL_CHECK(
+        cnnlGetCopyWorkspaceSize(cnnl_handle_.get(), input_desc_.get(),
+                                 out_desc_.get(), &workspace_size_));
+
+    default_workspace_ = cnrt_utils::AllocateDeviceBuffer(workspace_size_);
+  }
+
+  void operator()(const Tensor src, const bool /*non_blocking*/,
+                  Tensor out) const override {
+    if (output_size_ == 0) {
+      return;
+    }
+
+    INFINI_OPS_CNNL_CHECK(cnnlSetQueue(
+        cnnl_handle_.get(), static_cast<cnrtQueue_t>(stream_ ? stream_ : 0)));
+
+    void* workspace = workspace_ ? workspace_ : default_workspace_.get();
+    const auto workspace_size =
+        workspace_ ? workspace_size_in_bytes_ : workspace_size_;
+    assert(workspace_size >= workspace_size_ &&
+           "`CambriconCopy` requires a sufficiently large workspace.");
+
+    INFINI_OPS_CNNL_CHECK(cnnlCopy_v2(cnnl_handle_.get(), input_desc_.get(),
+                                      src.data(), out_desc_.get(), out.data(),
+                                      workspace, workspace_size_));
+  }
+
+  std::size_t workspace_size_in_bytes() const override {
+    return workspace_size_;
+  }
+
+ private:
+  std::size_t workspace_size_{0};
+
+  cnrt_utils::DeviceBuffer default_workspace_{};
+
+  cnnl_utils::Handle cnnl_handle_{};
+
+  cnnl_utils::TensorDescriptor input_desc_{};
+
+  cnnl_utils::TensorDescriptor out_desc_{};
+};
+
+}  // namespace infini::ops
+
+#endif

tests/test_copy.py

@@ -9,6 +9,7 @@
 @pytest.mark.parametrize(
     "input_shape, out_shape, input_strides, out_strides",
     (
+        ((), (), None, None),
         ((100, 100), (100, 100), (1, 100), (100, 1)),
         ((2, 2, 2, 4), (2, 2, 2, 4), (16, 8, 4, 1), (16, 8, 1, 2)),
         ((8, 4, 20, 64), (8, 4, 20, 64), (5120, 64, 256, 1), None),
@@ -59,3 +60,20 @@ def _torch_copy(input, non_blocking, out):
     out.copy_(input, non_blocking=non_blocking)
 
     return out
+
+
+@pytest.mark.parametrize("non_blocking", (False, True))
+@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16))
+def test_copy_preserves_kv_cache_outside_slice(dtype, non_blocking, device):
+    input = torch.randn((2, 3, 3, 4), dtype=dtype, device=device).transpose(1, 2)
+    cache = torch.full((2, 3, 11, 4), 17, dtype=dtype, device=device)
+    out = cache[:, :, 4:7, :]
+
+    before = cache.clone()
+    expected = input.expand_as(out).clone()
+
+    _copy(input, non_blocking, out)
+
+    assert torch.equal(out, expected)
+    assert torch.equal(cache[:, :, :4, :], before[:, :, :4, :])
+    assert torch.equal(cache[:, :, 7:, :], before[:, :, 7:, :])