// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <ATen/Functions.h>
#include <ATen/core/TensorBody.h>
#include <ATen/ops/arange.h>
#include <ATen/ops/empty_like.h>
#include <ATen/ops/eye.h>
#include <ATen/ops/full.h>
#include <ATen/ops/ones.h>
#include <ATen/ops/zeros.h>
#include <ATen/ops/zeros_like.h>

#include <optional>

#include "gtest/gtest.h"

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)

namespace {

void AssertPinned(const at::Tensor& t) {
  ASSERT_TRUE(t.is_pinned());
  ASSERT_FALSE(t.is_cuda());
}

void AssertNotPinned(const at::Tensor& t) { ASSERT_FALSE(t.is_pinned()); }

}  // namespace

TEST(ATenPinMemoryCreationTest, FullPinMemory) {
  // Test using TensorOptions with pinned_memory
  auto by_options = at::full(
      {2, 3}, 1.5f, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(by_options);

  // Test using explicit arguments with CPU device (should succeed)
  auto by_args =
      at::full({2, 3}, 1.5f, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(by_args);

  // Test without pin_memory
  auto no_pin =
      at::full({2, 3}, 1.5f, at::kFloat, std::nullopt, at::kCPU, false);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, FullPinMemoryWithCUDADeviceErrors) {
  // Test that pin_memory=true with explicit CUDA device throws error
  ASSERT_THROW(
      at::full({2, 3}, 1.5f, at::kFloat, std::nullopt, at::kCUDA, true),
      std::exception);
}

TEST(ATenPinMemoryCreationTest, OnesPinMemory) {
  auto by_options = at::ones(
      {4, 2}, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(by_options);

  auto by_args = at::ones({4, 2}, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(by_args);

  auto no_pin = at::ones({4, 2}, at::kFloat, std::nullopt, at::kCPU, false);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, OnesPinMemoryWithCUDADeviceErrors) {
  ASSERT_THROW(at::ones({4, 2}, at::kFloat, std::nullopt, at::kCUDA, true),
               std::exception);
}

TEST(ATenPinMemoryCreationTest, ZerosPinMemory) {
  auto by_options = at::zeros(
      {3, 5}, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(by_options);

  auto by_args = at::zeros({3, 5}, at::kFloat, at::kStrided, at::kCPU, true);
  AssertPinned(by_args);

  auto no_pin = at::zeros({3, 5}, at::kFloat, at::kStrided, at::kCPU, false);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, ZerosPinMemoryWithCUDADeviceErrors) {
  ASSERT_THROW(at::zeros({3, 5}, at::kFloat, at::kStrided, at::kCUDA, true),
               std::exception);
}

TEST(ATenPinMemoryCreationTest, EyePinMemory) {
  auto by_options =
      at::eye(6, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(by_options);

  auto by_args = at::eye(6, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(by_args);

  auto no_pin = at::eye(6, at::kFloat, std::nullopt, at::kCPU, false);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, EyePinMemoryWithCUDADeviceErrors) {
  ASSERT_THROW(at::eye(6, at::kFloat, std::nullopt, at::kCUDA, true),
               std::exception);
}

TEST(ATenPinMemoryCreationTest, ArangePinMemoryOverloads) {
  auto end_only_by_options =
      at::arange(10, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(end_only_by_options);

  auto end_only_by_args =
      at::arange(10, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(end_only_by_args);

  auto by_options = at::arange(
      0, 10, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(by_options);

  auto by_args = at::arange(0, 10, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(by_args);

  auto step_by_options = at::arange(
      0, 10, 2, at::TensorOptions().dtype(at::kFloat).pinned_memory(true));
  AssertPinned(step_by_options);

  auto step_by_args =
      at::arange(0, 10, 2, at::kFloat, std::nullopt, at::kCPU, true);
  AssertPinned(step_by_args);

  auto no_pin = at::arange(0, 10, at::kFloat, std::nullopt, at::kCPU, false);
  AssertNotPinned(no_pin);

  auto step_no_pin =
      at::arange(0, 10, 2, at::kFloat, std::nullopt, at::kCPU, false);
  AssertNotPinned(step_no_pin);
}

TEST(ATenPinMemoryCreationTest, ArangePinMemoryWithCUDADeviceErrors) {
  ASSERT_THROW(at::arange(10, at::kFloat, std::nullopt, at::kCUDA, true),
               std::exception);
  ASSERT_THROW(at::arange(0, 10, at::kFloat, std::nullopt, at::kCUDA, true),
               std::exception);
  ASSERT_THROW(at::arange(0, 10, 2, at::kFloat, std::nullopt, at::kCUDA, true),
               std::exception);
}

TEST(ATenPinMemoryCreationTest, EmptyLikePinMemory) {
  auto base = at::ones({2, 4}, at::kFloat);

  auto by_options =
      at::empty_like(base,
                     at::TensorOptions().dtype(at::kFloat).pinned_memory(true),
                     std::nullopt);
  AssertPinned(by_options);

  auto by_args = at::empty_like(
      base, at::kFloat, at::kStrided, at::kCPU, true, std::nullopt);
  AssertPinned(by_args);

  auto no_pin = at::empty_like(
      base, at::kFloat, at::kStrided, at::kCPU, false, std::nullopt);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, EmptyLikePinMemoryWithCUDADeviceErrors) {
  auto base = at::ones({2, 4}, at::kFloat);
  ASSERT_THROW(at::empty_like(base,
                              at::TensorOptions()
                                  .dtype(at::kFloat)
                                  .device(at::kCUDA)
                                  .pinned_memory(true),
                              std::nullopt),
               std::exception);
}

TEST(ATenPinMemoryCreationTest, ZerosLikePinMemory) {
  auto base = at::ones({2, 4}, at::kFloat);

  auto by_options =
      at::zeros_like(base,
                     at::TensorOptions().dtype(at::kFloat).pinned_memory(true),
                     std::nullopt);
  AssertPinned(by_options);

  auto by_args = at::zeros_like(
      base, at::kFloat, at::kStrided, at::kCPU, true, std::nullopt);
  AssertPinned(by_args);

  auto no_pin = at::zeros_like(
      base, at::kFloat, at::kStrided, at::kCPU, false, std::nullopt);
  AssertNotPinned(no_pin);
}

TEST(ATenPinMemoryCreationTest, ZerosLikePinMemoryWithCUDADeviceErrors) {
  auto base = at::ones({2, 4}, at::kFloat);
  ASSERT_THROW(at::zeros_like(base,
                              at::TensorOptions()
                                  .dtype(at::kFloat)
                                  .device(at::kCUDA)
                                  .pinned_memory(true),
                              std::nullopt),
               std::exception);
}

#endif  // PADDLE_WITH_CUDA || PADDLE_WITH_HIP