Refactor: refactor cg interface (#3293)

* init implementation

* fix cuda implementation

* address parallel problems

* fix compilation errors

* fix ut errors

* fix sdft ut error

* fix ut error

* fix typo

---------

Co-authored-by: Qianrui Liu <[email protected]>

Author: denghuilu

source/module_base/module_container/ATen/core/tensor.cpp

@@ -9,7 +9,7 @@ namespace container {
 
 Tensor::Tensor() : Tensor(DataType::DT_FLOAT) {}
 
-Tensor::Tensor(DataType data_type) : Tensor(data_type, TensorShape({})) {}
+Tensor::Tensor(DataType data_type) : Tensor(data_type, TensorShape({1})) {}
 
 // Constructor that creates a tensor with the given data type and shape using the default allocator.
 Tensor::Tensor(DataType data_type, const TensorShape& shape)
@@ -291,21 +291,31 @@ bool Tensor::AllocateFrom(const Tensor& other, const TensorShape& shape) {
 
 void Tensor::sync(const Tensor& rhs) {
     REQUIRES_OK(this->data_type_ == rhs.data_type_ 
-        && this->device_ == rhs.device_
-        && this->shape_ == rhs.shape_)
+        && this->device_ == rhs.device_)
 
-    TEMPLATE_ALL_2(data_type_, device_,
-            kernels::synchronize_memory<T_, DEVICE_, DEVICE_>()(
-                    this->data<T_>(), rhs.data<T_>(), this->NumElements()))
+    if (this->shape_ == rhs.shape_) {
+        TEMPLATE_ALL_2(data_type_, device_,
+                kernels::synchronize_memory<T_, DEVICE_, DEVICE_>()(
+                        this->data<T_>(), rhs.data<T_>(), this->NumElements()))
+    }
+    else {
+        TEMPLATE_ALL_2(data_type_, device_,
+                kernels::synchronize_memory_stride<T_, DEVICE_, DEVICE_>()(
+                        this->data<T_>(), rhs.data<T_>(), this->shape().dims(), rhs.shape().dims()))
+    }
 }
 
 Tensor Tensor::operator[](const int& index) const {
     REQUIRES_OK(
-        index > 0 && index < shape_.dim_size(0),
+        index >= 0 && index < shape_.dim_size(0),
         "Tensor index is out of bounds.")
 
     TensorShape output_shape = this->shape_;
     output_shape.remove_dim(0);
+    if (output_shape.ndim() == 0) {
+        // If the output shape is empty, we need to add a dimension of size 1
+        output_shape.add_dim(1);
+    }
     auto data_ = reinterpret_cast<char*>(this->data()) + index * shape_.strides()[0] * SizeOfType(this->data_type_);
 
     return TensorMap(data_, this->data_type_, this->device_, output_shape);

source/module_base/module_container/ATen/core/tensor.h

@@ -499,6 +499,12 @@ class Tensor {
         return this->NumElements() > 0;
     }
 
+    template<typename T>
+    void set_value(T value) {
+        TEMPLATE_ALL_2(this->data_type_, this->device_,
+            kernels::set_memory<T, DEVICE_>()(this->data<T>(), value, this->NumElements()))
+    }
+
 protected:
 
     /**

source/module_base/module_container/ATen/core/tensor_utils.h

@@ -270,7 +270,16 @@ void _internal_output(
 
 template <typename T>
 T extract(const container::Tensor& tensor) {
-    return reinterpret_cast<T*>(tensor.data())[0];
+    if (tensor.device_type() == DeviceType::CpuDevice) {
+        return reinterpret_cast<T*>(tensor.data())[0];
+    }
+    else {
+        T result = 0;
+        TEMPLATE_ALL_2(tensor.data_type(), tensor.device_type(),
+            kernels::synchronize_memory<T, DEVICE_CPU, DEVICE_>()(
+                &result, reinterpret_cast<T*>(tensor.data()), 1))
+        return result;
+    }
 }
 
 } // namespace container

source/module_base/module_container/ATen/kernels/memory.h

@@ -4,6 +4,7 @@
 #include <vector>
 #include <complex>
 
+#include <base/macros/macros.h>
 #include <ATen/core/tensor_types.h>
 
 namespace container {
@@ -72,6 +73,26 @@ struct synchronize_memory {
         const size_t& size);
 };
 
+template <typename T, typename Device_out, typename Device_in>
+struct synchronize_memory_stride {
+    void operator()(
+        T* arr_out,
+        const T* arr_in,
+        const std::vector<int64_t>& out_size,
+        const std::vector<int64_t>& in_size)
+    {
+        REQUIRES_OK(in_size.size() == out_size.size() && in_size.size() <= 2);
+        if (in_size.size() == 1) {
+            synchronize_memory<T, Device_out, Device_in>()(arr_out, arr_in, in_size[0]);
+        }
+        else {
+            for (int64_t ii = 0; ii < out_size[0]; ii++) {
+                synchronize_memory<T, Device_out, Device_in>()(arr_out + ii * out_size[1], arr_in + ii * in_size[1], in_size[1]);
+            }
+        }
+    }
+};
+
 /**
  * @brief Casts memory between devices.
  *

source/module_hsolver/diagh.h

@@ -28,7 +28,9 @@ class DiagH
     // virtual void init()=0;
     std::string method = "none";
 
-    virtual void diag(hamilt::Hamilt<T, Device> *phm_in, psi::Psi<T, Device> &psi, Real *eigenvalue_in) = 0;
+    virtual void diag(hamilt::Hamilt<T, Device> *phm_in, psi::Psi<T, Device> &psi, Real *eigenvalue_in) {
+      ModuleBase::WARNING_QUIT("diagh", "diag method not implemented for the base class!");
+    };
 
 };