diff --git a/src/native/cambricon/cnnl_utils.h b/src/native/cambricon/cnnl_utils.h new file mode 100644 index 000000000..e5c78754c --- /dev/null +++ b/src/native/cambricon/cnnl_utils.h @@ -0,0 +1,97 @@ +#ifndef INFINI_OPS_CAMBRICON_CNNL_UTILS_H_ +#define INFINI_OPS_CAMBRICON_CNNL_UTILS_H_ + +#include +#include +#include +#include + +#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 { + using pointer = cnnlHandle_t; + + void operator()(pointer handle) const noexcept { + if (handle) { + (void)cnnlDestroy(handle); + } + } +}; + +using Handle = + std::unique_ptr, HandleDeleter>; + +inline Handle CreateHandle() { + cnnlHandle_t handle{nullptr}; + INFINI_OPS_CNNL_CHECK(cnnlCreate(&handle)); + + return Handle{handle}; +} + +struct TensorDescriptorDeleter { + using pointer = cnnlTensorDescriptor_t; + + void operator()(pointer desc) const noexcept { + if (desc) { + (void)cnnlDestroyTensorDescriptor(desc); + } + } +}; + +using TensorDescriptor = + std::unique_ptr, + TensorDescriptorDeleter>; + +inline TensorDescriptor CreateTensorDescriptor() { + cnnlTensorDescriptor_t desc{nullptr}; + INFINI_OPS_CNNL_CHECK(cnnlCreateTensorDescriptor(&desc)); + + return TensorDescriptor{desc}; +} + +inline void SetTensorDescriptor(cnnlTensorDescriptor_t desc, DataType dtype, + const Tensor::Shape& shape, + const Tensor::Strides& 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 = + shape.empty() ? std::vector{1} + : std::vector(shape.begin(), shape.end()); + const auto cnnl_strides = + strides.empty() + ? std::vector{1} + : std::vector(strides.begin(), strides.end()); + const auto ndim = static_cast(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; +} + +} // namespace infini::ops::cnnl_utils + +#endif diff --git a/src/native/cambricon/cnrt_utils.h b/src/native/cambricon/cnrt_utils.h new file mode 100644 index 000000000..6bd55ce7a --- /dev/null +++ b/src/native/cambricon/cnrt_utils.h @@ -0,0 +1,38 @@ +#ifndef INFINI_OPS_CAMBRICON_CNRT_UTILS_H_ +#define INFINI_OPS_CAMBRICON_CNRT_UTILS_H_ + +#include + +#include +#include +#include + +namespace infini::ops::cnrt_utils { + +struct DeviceBufferDeleter { + using pointer = void*; + + void operator()(pointer buffer) const noexcept { + if (buffer) { + (void)cnrtFree(buffer); + } + } +}; + +using DeviceBuffer = std::unique_ptr; + +inline DeviceBuffer AllocateDeviceBuffer(std::size_t size) { + if (size == 0) { + return {}; + } + + void* buffer{nullptr}; + [[maybe_unused]] const auto status = cnrtMalloc(&buffer, size); + assert(status == cnrtSuccess && "`cnrtMalloc` failed."); + + return DeviceBuffer{buffer}; +} + +} // namespace infini::ops::cnrt_utils + +#endif diff --git a/src/native/cambricon/ops/copy/cnnl.h b/src/native/cambricon/ops/copy/cnnl.h new file mode 100644 index 000000000..23dde7a66 --- /dev/null +++ b/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 +#include + +#include "base/copy.h" +#include "native/cambricon/cnnl_utils.h" +#include "native/cambricon/cnrt_utils.h" + +namespace infini::ops { + +template <> +class Operator : 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(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 diff --git a/tests/test_copy.py b/tests/test_copy.py index 61bacfa11..126ed1e39 100644 --- a/tests/test_copy.py +++ b/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:, :])