#include "operator.h"

#include "source_base/timer.h"

#include "source_base/tool_quit.h"

using namespace hamilt;

template <typename T, typename Device>

Operator<T, Device>::Operator()

{

}

template <typename T, typename Device>

Operator<T, Device>::~Operator()

{

if (this->hpsi != nullptr)

{

delete this->hpsi;

}

Operator* last = this->next_op;

Operator* last_sub = this->next_sub_op;

while (last != nullptr || last_sub != nullptr)

{

if (last_sub != nullptr)

{ // delete sub_chain first

Operator* node_delete = last_sub;

last_sub = last_sub->next_sub_op;

node_delete->next_sub_op = nullptr;

delete node_delete;

}

else

{ // delete main chain if sub_chain is deleted

Operator* node_delete = last;

last_sub = last->next_sub_op;

node_delete->next_sub_op = nullptr;

last = last->next_op;

node_delete->next_op = nullptr;

delete node_delete;

}

}

}

template <typename T, typename Device>

typename Operator<T, Device>::hpsi_info Operator<T, Device>::hPsi(hpsi_info& input) const

{

using syncmem_op = base_device::memory::synchronize_memory_op<T, Device, Device>;

auto psi_input = std::get<0>(input);

std::tuple<const T*, int> psi_info = psi_input->to_range(std::get<1>(input));

int nbands = std::get<1>(psi_info);

T* tmhpsi = this->get_hpsi(input);

const T* tmpsi_in = std::get<0>(psi_info);

// if range in hpsi_info is illegal, the first return of to_range() would be nullptr

if (tmpsi_in == nullptr)

{

ModuleBase::WARNING_QUIT("Operator", "please choose correct range of psi for hPsi()!");

}

// if in_place, copy temporary hpsi to target hpsi_pointer, then delete hpsi and new a wrapper for return

T* hpsi_pointer = std::get<2>(input);

if (this->in_place)

{

syncmem_op()(hpsi_pointer, this->hpsi->get_pointer(), this->hpsi->size());

delete this->hpsi;

this->hpsi = new psi::Psi<T, Device>(hpsi_pointer,

1,

nbands / psi_input->get_npol(),

psi_input->get_nbasis(),

psi_input->get_nbasis(),

true);

}

auto call_act = [&, this](const Operator* op, const bool& is_first_node) -> void {

// a "psi" with the bands of needed range

psi::Psi<T, Device> psi_wrapper(const_cast<T*>(tmpsi_in),

1,

nbands,

psi_input->get_nbasis(),

psi_input->get_nbasis(),

true);

switch (op->get_act_type())

{

case 2:

op->act(psi_wrapper, *this->hpsi, nbands);

break;

default:

op->act(nbands,

psi_input->get_nbasis(),

psi_input->get_npol(),

tmpsi_in,

this->hpsi->get_pointer(),

psi_input->get_current_nbas(),

is_first_node);

break;

}

};

ModuleBase::timer::start("Operator", "hPsi");

call_act(this, true); // first node

Operator* node((Operator*)this->next_op);

while (node != nullptr)

{

call_act(node, false); // other nodes

node = (Operator*)(node->next_op);

}

ModuleBase::timer::end("Operator", "hPsi");

return hpsi_info(this->hpsi, psi::Range(1, 0, 0, nbands / psi_input->get_npol()), hpsi_pointer);

}

template <typename T, typename Device>

void Operator<T, Device>::init(const int ik_in)

{

this->ik = ik_in;

if (this->next_op != nullptr)

{

this->next_op->init(ik_in);

}

}

template <typename T, typename Device>

void Operator<T, Device>::add(Operator* next)

{

if (next == nullptr)

{

return;

}

next->is_first_node = false;

if (next->next_op != nullptr)

{

this->add(next->next_op);

}

Operator* last = this;

// loop to end of the chain

while (last->next_op != nullptr)

{

if (next->cal_type == last->cal_type)

{

break;

}

last = last->next_op;

}

if (next->cal_type == last->cal_type)

{

// insert next to sub chain of current node

Operator* sub_last = last;

while (sub_last->next_sub_op != nullptr)

{

sub_last = sub_last->next_sub_op;

}

sub_last->next_sub_op = next;

return;

}

else

{

last->next_op = next;

}

}

template <typename T, typename Device>

T* Operator<T, Device>::get_hpsi(const hpsi_info& info) const

{

const int nbands_range = (std::get<1>(info).range_2 - std::get<1>(info).range_1 + 1);

// in_place call of hPsi, hpsi inputs as new psi,

// create a new hpsi and delete old hpsi later

T* hpsi_pointer = std::get<2>(info);

const T* psi_pointer = std::get<0>(info)->get_pointer();

if (this->hpsi != nullptr)

{

delete this->hpsi;

this->hpsi = nullptr;

}

if (!hpsi_pointer)

{

ModuleBase::WARNING_QUIT("Operator::hPsi", "hpsi_pointer can not be nullptr");

}

else if (hpsi_pointer == psi_pointer)

{

this->in_place = true;

// this->hpsi = new psi::Psi<T, Device>(std::get<0>(info)[0], 1, nbands_range);

this->hpsi = new psi::Psi<T, Device>(1,

nbands_range,

std::get<0>(info)->get_nbasis(),

std::get<0>(info)->get_nbasis(),

true);

}

else

{

this->in_place = false;

this->hpsi = new psi::Psi<T, Device>(hpsi_pointer,

1,

nbands_range,

std::get<0>(info)->get_nbasis(),

std::get<0>(info)->get_nbasis(),

true);

}

hpsi_pointer = this->hpsi->get_pointer();

size_t total_hpsi_size = nbands_range * this->hpsi->get_nbasis();

// ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size);

// denghui replaced at 20221104

// set_memory_op()(this->ctx, hpsi_pointer, 0, total_hpsi_size);

return hpsi_pointer;

}

namespace hamilt

{

template class Operator<float, base_device::DEVICE_CPU>;

template class Operator<std::complex<float>, base_device::DEVICE_CPU>;

template class Operator<double, base_device::DEVICE_CPU>;

template class Operator<std::complex<double>, base_device::DEVICE_CPU>;

#if ((defined __CUDA) || (defined __ROCM))

template class Operator<float, base_device::DEVICE_GPU>;

template class Operator<std::complex<float>, base_device::DEVICE_GPU>;

template class Operator<double, base_device::DEVICE_GPU>;

template class Operator<std::complex<double>, base_device::DEVICE_GPU>;

#endif

} // namespace hamilt