diff --git a/src/native/cambricon/cnnl_utils.h b/src/native/cambricon/cnnl_utils.h index a32c65f10..e5c78754c 100644 --- a/src/native/cambricon/cnnl_utils.h +++ b/src/native/cambricon/cnnl_utils.h @@ -3,14 +3,19 @@ #include #include -#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 { @@ -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,68 +54,33 @@ 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 -int CheckedInt(Integer value) { - static_assert(std::is_integral_v); - - [[maybe_unused]] bool out_of_range{false}; - if constexpr (std::is_signed_v) { - const auto wide = static_cast(value); - out_of_range = wide < std::numeric_limits::min() || - wide > std::numeric_limits::max(); - } else { - const auto wide = static_cast(value); - out_of_range = - wide > static_cast(std::numeric_limits::max()); - } - - assert(!out_of_range && - "`CNNL tensor descriptor` value does not fit in `int`."); - - return static_cast(value); -} - -template -std::vector CheckedIntVector(const Values& values) { - std::vector 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{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, @@ -119,6 +88,7 @@ inline TensorDescriptor MakeTensorDescriptor(DataType dtype, const Tensor::Strides& strides) { auto desc = CreateTensorDescriptor(); SetTensorDescriptor(desc.get(), dtype, shape, strides); + return desc; } 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:, :])