Skip to content
Open
16 changes: 10 additions & 6 deletions source/source_base/module_container/ATen/core/tensor_accessor.h
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
#ifndef ATEN_CORE_TENSOR_ACCESSOR_H_
#define ATEN_CORE_TENSOR_ACCESSOR_H_

#include <cstddef> // Include the <cstddef> header file to define size_t
#include <cstddef>
#include <cstdint>
#include <base/macros/macros.h>
#include <base/utils/array_ref.h>

namespace container {

Expand Down Expand Up @@ -33,12 +32,17 @@ class TensorAccessorBase {
const index_t* strides)
: data_(data), sizes_(sizes), strides_(strides) {}

AT_HOST int_array_ref sizes() const {
return {sizes_, N};
// NOTE: Previously returned int_array_ref (from <base/utils/array_ref.h>),
// which supported .size(), iteration, etc. Now returns raw const pointer.
// If you need array_ref functionality, include <base/utils/array_ref.h> directly
// and wrap the result: int_array_ref(accessor.sizes(), N).
// mohan add 20260605
AT_HOST_DEVICE const index_t* sizes() const {
return sizes_;
}

AT_HOST int_array_ref strides() const {
return {strides_, N};
AT_HOST_DEVICE const index_t* strides() const {
return strides_;
}

AT_HOST_DEVICE index_t stride(index_t idx) const {
Expand Down
96 changes: 96 additions & 0 deletions source/source_base/module_container/ATen/core/tensor_enums.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Extracted from tensor_types.h to break the circular dependency:
// macros.h -> tensor_types.h -> cuda.h/rocm.h -> macros.h
// This file provides only enums and template specializations needed by macros.h,
// without any GPU library dependencies.
// mohan add 20260605
#ifndef ATEN_CORE_TENSOR_ENUMS_H_
#define ATEN_CORE_TENSOR_ENUMS_H_

#include <complex>

#include "source_base/module_device/types.h"

namespace container {

enum class DataType {
DT_INVALID = 0,
DT_FLOAT = 1,
DT_DOUBLE = 2,
DT_INT = 3,
DT_INT64 = 4,
DT_COMPLEX = 5,
DT_COMPLEX_DOUBLE = 6,
};

struct DEVICE_CPU {};
struct DEVICE_GPU {};

enum class DeviceType {
UnKnown = 0,
CpuDevice = 1,
GpuDevice = 2,
};

template <typename T>
struct DeviceTypeToEnum {
static constexpr DeviceType value = {};
};

template <>
struct DeviceTypeToEnum<DEVICE_CPU> {
static constexpr DeviceType value = DeviceType::CpuDevice;
};

template <>
struct DeviceTypeToEnum<DEVICE_GPU> {
static constexpr DeviceType value = DeviceType::GpuDevice;
};

template <>
struct DeviceTypeToEnum<base_device::DEVICE_CPU> {
static constexpr DeviceType value = DeviceType::CpuDevice;
};

template <>
struct DeviceTypeToEnum<base_device::DEVICE_GPU> {
static constexpr DeviceType value = DeviceType::GpuDevice;
};

template <typename T>
struct DataTypeToEnum {
static constexpr DataType value = {};
};

template <>
struct DataTypeToEnum<int> {
static constexpr DataType value = DataType::DT_INT;
};

template <>
struct DataTypeToEnum<float> {
static constexpr DataType value = DataType::DT_FLOAT;
};

template <>
struct DataTypeToEnum<double> {
static constexpr DataType value = DataType::DT_DOUBLE;
};

template <>
struct DataTypeToEnum<int64_t> {
static constexpr DataType value = DataType::DT_INT64;
};

template <>
struct DataTypeToEnum<std::complex<float>> {
static constexpr DataType value = DataType::DT_COMPLEX;
};

template <>
struct DataTypeToEnum<std::complex<double>> {
static constexpr DataType value = DataType::DT_COMPLEX_DOUBLE;
};

} // namespace container

#endif // ATEN_CORE_TENSOR_ENUMS_H_
139 changes: 11 additions & 128 deletions source/source_base/module_container/ATen/core/tensor_types.h
Original file line number Diff line number Diff line change
Expand Up @@ -5,26 +5,23 @@
#ifndef ATEN_CORE_TENSOR_TYPES_H_
#define ATEN_CORE_TENSOR_TYPES_H_

#include <regex>
#include <string>
#include <vector>
#include <numeric>
#include <complex>
#include <utility>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <algorithm>
#include <unordered_map>
#include <initializer_list>

#include "source_base/module_device/types.h"

#if defined(__CUDACC__)
#include <base/macros/cuda.h>
#elif defined(__HIPCC__)
#include <base/macros/rocm.h>
#endif // defined(__CUDACC__) || defined(__HIPCC__)
#include "ATen/core/tensor_enums.h"

// NOTE: Previously this file included <base/macros/cuda.h> or <base/macros/rocm.h>,
// which transitively provided GetTypeThrust, GetTypeCuda, GetTypeRocm, etc.
// Now replaced with <thrust/complex.h> which only provides thrust::complex type.
// If you need GetTypeThrust/GetTypeCuda/GetTypeRocm, include <base/macros/cuda.h>
// or <base/macros/rocm.h> directly in your .cpp file.
// mohan add 20260605
#if defined(__CUDACC__) || defined(__HIPCC__)
#include <thrust/complex.h>
#endif

namespace container {

Expand All @@ -41,41 +38,6 @@ static inline bool element_compare(T& a, T& b) {
}
}

/**
@brief Enumeration of data types for tensors.
The DataType enum lists the supported data types for tensors. Each data type
is identified by a unique value. The DT_INVALID value is reserved for invalid
data types.
*/
enum class DataType {
DT_INVALID = 0, ///< Invalid data type */
DT_FLOAT = 1, ///< Single-precision floating point */
DT_DOUBLE = 2, ///< Double-precision floating point */
DT_INT = 3, ///< 32-bit integer */
DT_INT64 = 4, ///< 64-bit integer */
DT_COMPLEX = 5, ///< 32-bit complex */
DT_COMPLEX_DOUBLE = 6, /**< 64-bit complex */
// ... other data types
};

/**
*@struct DEVICE_CPU, DEVICE_GPU
*@brief A tag type for identifying CPU and GPU devices.
*/
struct DEVICE_CPU;
struct DEVICE_GPU;

struct DEVICE_CPU {};
struct DEVICE_GPU {};
/**
* @brief The type of memory used by an allocator.
*/
enum class DeviceType {
UnKnown = 0, ///< Memory type is unknown.
CpuDevice = 1, ///< Memory type is CPU.
GpuDevice = 2, ///< Memory type is GPU(CUDA or ROCm).
};

/**
* @brief Template struct to determine the return type based on the input type.
*
Expand Down Expand Up @@ -141,85 +103,6 @@ struct ContainerToPsi<container::DEVICE_GPU> {
using type = base_device::DEVICE_GPU; /**< The return type specialization for std::complex<double>. */
};


/**
* @brief Template struct for mapping a Device Type to its corresponding enum value.
*
* @param T The DataType to map to its enum value.
*
* @return The enumeration value corresponding to the data type.
* This method uses template specialization to map each supported data type to its
* corresponding enumeration value. If the template argument T is not a supported
* data type, this method will cause a compile-time error.
* Example usage:
* DataTypeToEnum<float>::value; // Returns DataType::DT_FLOAT
*/
template <typename T>
struct DeviceTypeToEnum {
static constexpr DeviceType value = {};
};
// Specializations of DeviceTypeToEnum for supported devices.
template <>
struct DeviceTypeToEnum<DEVICE_CPU> {
static constexpr DeviceType value = DeviceType::CpuDevice;
};
template <>
struct DeviceTypeToEnum<DEVICE_GPU> {
static constexpr DeviceType value = DeviceType::GpuDevice;
};
template <>
struct DeviceTypeToEnum<base_device::DEVICE_CPU>
{
static constexpr DeviceType value = DeviceType::CpuDevice;
};
template <>
struct DeviceTypeToEnum<base_device::DEVICE_GPU>
{
static constexpr DeviceType value = DeviceType::GpuDevice;
};

/**
* @brief Template struct for mapping a DataType to its corresponding enum value.
*
* @param T The DataType to map to its enum value.
*
* @return The enumeration value corresponding to the data type.
* This method uses template specialization to map each supported data type to its
* corresponding enumeration value. If the template argument T is not a supported
* data type, this method will cause a compile-time error.
* Example usage:
* DataTypeToEnum<float>::value; // Returns DataType::DT_FLOAT
*/
template <typename T>
struct DataTypeToEnum {
static constexpr DataType value = {};
};
// Specializations of DataTypeToEnum for supported types.
template <>
struct DataTypeToEnum<int> {
static constexpr DataType value = DataType::DT_INT;
};
template <>
struct DataTypeToEnum<float> {
static constexpr DataType value = DataType::DT_FLOAT;
};
template <>
struct DataTypeToEnum<double> {
static constexpr DataType value = DataType::DT_DOUBLE;
};
template <>
struct DataTypeToEnum<int64_t> {
static constexpr DataType value = DataType::DT_INT64;
};
template <>
struct DataTypeToEnum<std::complex<float>> {
static constexpr DataType value = DataType::DT_COMPLEX;
};
template <>
struct DataTypeToEnum<std::complex<double>> {
static constexpr DataType value = DataType::DT_COMPLEX_DOUBLE;
};

#if defined(__CUDACC__) || defined(__HIPCC__)
template <>
struct DataTypeToEnum<thrust::complex<float>> {
Expand Down
2 changes: 2 additions & 0 deletions source/source_base/module_container/ATen/core/tensor_utils.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#ifndef ATEN_CORE_TENSOR_UTILS_H_
#define ATEN_CORE_TENSOR_UTILS_H_

#include <iomanip>

#include <ATen/core/tensor.h>
#include <ATen/core/tensor_shape.h>

Expand Down
3 changes: 3 additions & 0 deletions source/source_base/module_container/ATen/ops/einsum_op.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

#include <array>
#include <algorithm>
#include <functional>
#include <numeric>
#include <regex>

#include <ATen/ops/linalg_op.h>
#include <ATen/core/tensor_types.h>
Expand Down
2 changes: 2 additions & 0 deletions source/source_base/module_container/ATen/ops/einsum_op.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#ifndef ATEN_KERNELS_EINSUM_OP_H_
#define ATEN_KERNELS_EINSUM_OP_H_

#include <unordered_map>

#include <ATen/core/tensor.h>
#include <ATen/core/tensor_map.h>

Expand Down
9 changes: 8 additions & 1 deletion source/source_base/module_container/base/macros/macros.h
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,14 @@
#include <base/macros/rocm.h>
#endif

#include <ATen/core/tensor_types.h>
// NOTE: Previously included <ATen/core/tensor_types.h> which pulled in
// <base/macros/cuda.h> or <base/macros/rocm.h> transitively.
// Now only includes <ATen/core/tensor_enums.h> which provides DataType,
// DeviceType, DataTypeToEnum, DeviceTypeToEnum without GPU dependencies.
// If you need GetTypeReal, PsiToContainer, ContainerToPsi, etc.,
// include <ATen/core/tensor_types.h> directly.
// mohan add 20260605
#include <ATen/core/tensor_enums.h>

#define DISALLOW_COPY_AND_ASSIGN(TypeName) \
TypeName(const TypeName&) = delete; \
Expand Down
1 change: 1 addition & 0 deletions source/source_base/module_device/output_device.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#include <base/macros/macros.h>
#include <cstring>
#include <iostream>
#include <numeric>
#include <set>
#ifdef __MPI
#include "mpi.h"
Expand Down
1 change: 1 addition & 0 deletions source/source_esolver/esolver_ks_lcao.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_lcao/LCAO_domain.h" // need DeePKS_init
#include "source_lcao/FORCE_STRESS.h"
#include "source_lcao/module_gint/gint.h"
#include "source_estate/elecstate_lcao.h"
#include "source_lcao/hamilt_lcao.h"
#include "source_hsolver/hsolver_lcao.h"
Expand Down
1 change: 0 additions & 1 deletion source/source_esolver/esolver_ks_lcao.h
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
#include "esolver_ks.h"
#include "source_lcao/record_adj.h" // adjacent atoms
#include "source_basis/module_nao/two_center_bundle.h" // nao basis
#include "source_lcao/module_gint/gint.h" // gint
#include "source_lcao/module_gint/gint_info.h"
#include "source_estate/module_charge/gint_precision_controller.h"
#include "source_lcao/setup_deepks.h" // for deepks, mohan add 20251008
Expand Down
1 change: 1 addition & 0 deletions source/source_esolver/lcao_others.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_lcao/hamilt_lcao.h"
#include "source_lcao/module_dftu/dftu.h"
#include "source_lcao/module_gint/gint.h"
#include "source_base/formatter.h"
#include "source_base/timer.h"
#include "source_cell/module_neighbor/sltk_atom_arrange.h"
Expand Down
2 changes: 0 additions & 2 deletions source/source_estate/module_pot/H_TDDFT_pw.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,7 @@
#include "source_base/constants.h"
#include "source_base/math_integral.h"
#include "source_base/timer.h"
#include "source_io/module_parameter/input_conv.h"
#include "source_io/module_parameter/parameter.h"
#include "source_lcao/module_rt/evolve_elec.h"

namespace elecstate
{
Expand Down
2 changes: 2 additions & 0 deletions source/source_hsolver/diago_bpcg.h
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
#ifndef DIAGO_BPCG_H_
#define DIAGO_BPCG_H_

#include <functional>

#include "source_base/kernels/math_kernel_op.h"
#include "source_base/module_device/memory_op.h"
#include "source_base/module_device/types.h"
Expand Down
Loading
Loading