From 0ffbc5eb686317dfdcdd8e42b7e7afa6f3aa7aca Mon Sep 17 00:00:00 2001 From: windy-pig Date: Thu, 12 Jun 2025 18:42:54 +0800 Subject: [PATCH] Add CPU support for kernels. Modified `_hamiltonian_cuda.cu` to support CPU as `_hamiltonian_cpu.cpp`. The changes have been registered in `hamiltonian.py`. --- qmb/_hamiltonian_cpu.cpp | 593 +++++++++++++++++++++++++++++++++++++++ qmb/hamiltonian.py | 1 + 2 files changed, 594 insertions(+) create mode 100644 qmb/_hamiltonian_cpu.cpp diff --git a/qmb/_hamiltonian_cpu.cpp b/qmb/_hamiltonian_cpu.cpp new file mode 100644 index 0000000..7c59d94 --- /dev/null +++ b/qmb/_hamiltonian_cpu.cpp @@ -0,0 +1,593 @@ +#include +#include +#include + +namespace qmb_hamiltonian_cpu { + +constexpr torch::DeviceType device = torch::kCPU; + +template +struct array_less { + bool operator()(const std::array& lhs, const std::array& rhs) const { + for (std::int64_t i = 0; i < size; ++i) { + if (lhs[i] < rhs[i]) { + return true; + } + if (lhs[i] > rhs[i]) { + return false; + } + } + return false; + } +}; + +template +struct array_square_greater { + T square(const std::array& value) const { + T result = 0; + for (std::int64_t i = 0; i < size; ++i) { + result += value[i] * value[i]; + } + return result; + } + bool operator()(const std::array& lhs, const std::array& rhs) const { + return square(lhs) > square(rhs); + } +}; + +bool get_bit(std::uint8_t* data, std::uint8_t index) { + return ((*data) >> index) & 1; +} + +void set_bit(std::uint8_t* data, std::uint8_t index, bool value) { + if (value) { + *data |= (1 << index); + } else { + *data &= ~(1 << index); + } +} + +template +std::pair hamiltonian_apply_kernel( + std::array& current_configs, + std::int64_t term_index, + std::int64_t batch_index, + const std::array* site, // term_number + const std::array* kind // term_number +) { + static_assert(particle_cut == 1 || particle_cut == 2, "particle_cut != 1 or 2 not implemented"); + bool success = true; + bool parity = false; + for (std::int64_t op_index = max_op_number; op_index-- > 0;) { + std::int16_t site_single = site[term_index][op_index]; + std::uint8_t kind_single = kind[term_index][op_index]; + if (kind_single == 2) { + continue; + } + std::uint8_t to_what = kind_single; + if (get_bit(¤t_configs[site_single / 8], site_single % 8) == to_what) { + success = false; + break; + } + set_bit(¤t_configs[site_single / 8], site_single % 8, to_what); + if constexpr (particle_cut == 1) { + for (std::int16_t s = 0; s < site_single; ++s) { + parity ^= get_bit(¤t_configs[s / 8], s % 8); + } + } + } + return std::make_pair(success, parity); +} + +template +void apply_within_kernel( + std::int64_t term_index, + std::int64_t batch_index, + std::int64_t term_number, + std::int64_t batch_size, + std::int64_t result_batch_size, + const std::array* site, // term_number + const std::array* kind, // term_number + const std::array* coef, // term_number + const std::array* configs, // batch_size + const std::array* psi, // batch_size + const std::array* result_configs, // result_batch_size + std::array* result_psi +) { + std::array current_configs = configs[batch_index]; + auto [success, parity] = hamiltonian_apply_kernel( + /*current_configs=*/current_configs, + /*term_index=*/term_index, + /*batch_index=*/batch_index, + /*site=*/site, + /*kind=*/kind + ); + + if (!success) { + return; + } + success = false; + std::int64_t low = 0; + std::int64_t high = result_batch_size - 1; + std::int64_t mid = 0; + auto compare = array_less(); + while (low <= high) { + mid = (low + high) / 2; + if (compare(current_configs, result_configs[mid])) { + high = mid - 1; + } else if (compare(result_configs[mid], current_configs)) { + low = mid + 1; + } else { + success = true; + break; + } + } + if (!success) { + return; + } + std::int8_t sign = parity ? -1 : +1; + result_psi[mid][0] += sign * (coef[term_index][0] * psi[batch_index][0] - coef[term_index][1] * psi[batch_index][1]); + result_psi[mid][1] += sign * (coef[term_index][0] * psi[batch_index][1] + coef[term_index][1] * psi[batch_index][0]); +} + +template +void apply_within_kernel_interface( + std::int64_t term_number, + std::int64_t batch_size, + std::int64_t result_batch_size, + const std::array* site, // term_number + const std::array* kind, // term_number + const std::array* coef, // term_number + const std::array* configs, // batch_size + const std::array* psi, // batch_size + const std::array* result_configs, // result_batch_size + std::array* result_psi +) { + for (std::int64_t term_index = 0; term_index < term_number; ++term_index) { + for (std::int64_t batch_index = 0; batch_index < batch_size; ++batch_index) { + apply_within_kernel( + /*term_index=*/term_index, + /*batch_index=*/batch_index, + /*term_number=*/term_number, + /*batch_size=*/batch_size, + /*result_batch_size=*/result_batch_size, + /*site=*/site, + /*kind=*/kind, + /*coef=*/coef, + /*configs=*/configs, + /*psi=*/psi, + /*result_configs=*/result_configs, + /*result_psi=*/result_psi + ); + } + } +} + +template +auto apply_within_interface( + const torch::Tensor& configs, + const torch::Tensor& psi, + const torch::Tensor& result_configs, + const torch::Tensor& site, + const torch::Tensor& kind, + const torch::Tensor& coef +) -> torch::Tensor { + std::int64_t device_id = configs.device().index(); + std::int64_t batch_size = configs.size(0); + std::int64_t result_batch_size = result_configs.size(0); + std::int64_t term_number = site.size(0); + + TORCH_CHECK(configs.device().type() == torch::kCPU, "configs must be on CPU.") + TORCH_CHECK(configs.device().index() == device_id, "configs must be on the same device as others."); + TORCH_CHECK(configs.is_contiguous(), "configs must be contiguous.") + TORCH_CHECK(configs.dtype() == torch::kUInt8, "configs must be uint8.") + TORCH_CHECK(configs.dim() == 2, "configs must be 2D.") + TORCH_CHECK(configs.size(0) == batch_size, "configs batch size must match the provided batch_size."); + TORCH_CHECK(configs.size(1) == n_qubytes, "configs must have the same number of qubits as the provided n_qubytes."); + + TORCH_CHECK(psi.device().type() == torch::kCPU, "psi must be on CPU.") + TORCH_CHECK(psi.device().index() == device_id, "psi must be on the same device as others."); + TORCH_CHECK(psi.is_contiguous(), "psi must be contiguous.") + TORCH_CHECK(psi.dtype() == torch::kFloat64, "psi must be float64.") + TORCH_CHECK(psi.dim() == 2, "psi must be 2D.") + TORCH_CHECK(psi.size(0) == batch_size, "psi batch size must match the provided batch_size."); + TORCH_CHECK(psi.size(1) == 2, "psi must contain 2 elements for each batch."); + + TORCH_CHECK(result_configs.device().type() == torch::kCPU, "result_configs must be on CPU.") + TORCH_CHECK(result_configs.device().index() == device_id, "result_configs must be on the same device as others."); + TORCH_CHECK(result_configs.is_contiguous(), "result_configs must be contiguous.") + TORCH_CHECK(result_configs.dtype() == torch::kUInt8, "result_configs must be uint8.") + TORCH_CHECK(result_configs.dim() == 2, "result_configs must be 2D.") + TORCH_CHECK(result_configs.size(0) == result_batch_size, "result_configs batch size must match the provided result_batch_size.") + TORCH_CHECK(result_configs.size(1) == n_qubytes, "result_configs must have the same number of qubits as the provided n_qubytes."); + + TORCH_CHECK(site.device().type() == torch::kCPU, "site must be on CPU.") + TORCH_CHECK(site.device().index() == device_id, "site must be on the same device as others."); + TORCH_CHECK(site.is_contiguous(), "site must be contiguous.") + TORCH_CHECK(site.dtype() == torch::kInt16, "site must be int16.") + TORCH_CHECK(site.dim() == 2, "site must be 2D.") + TORCH_CHECK(site.size(0) == term_number, "site size must match the provided term_number."); + TORCH_CHECK(site.size(1) == max_op_number, "site must match the provided max_op_number."); + + TORCH_CHECK(kind.device().type() == torch::kCPU, "kind must be on CPU.") + TORCH_CHECK(kind.device().index() == device_id, "kind must be on the same device as others."); + TORCH_CHECK(kind.is_contiguous(), "kind must be contiguous.") + TORCH_CHECK(kind.dtype() == torch::kUInt8, "kind must be uint8.") + TORCH_CHECK(kind.dim() == 2, "kind must be 2D.") + TORCH_CHECK(kind.size(0) == term_number, "kind size must match the provided term_number."); + TORCH_CHECK(kind.size(1) == max_op_number, "kind must match the provided max_op_number."); + + TORCH_CHECK(coef.device().type() == torch::kCPU, "coef must be on CPU.") + TORCH_CHECK(coef.device().index() == device_id, "coef must be on the same device as others."); + TORCH_CHECK(coef.is_contiguous(), "coef must be contiguous.") + TORCH_CHECK(coef.dtype() == torch::kFloat64, "coef must be float64.") + TORCH_CHECK(coef.dim() == 2, "coef must be 2D.") + TORCH_CHECK(coef.size(0) == term_number, "coef size must match the provided term_number."); + TORCH_CHECK(coef.size(1) == 2, "coef must contain 2 elements for each term."); + + auto sorted_result_configs = result_configs.clone(torch::MemoryFormat::Contiguous); + auto result_sort_index = torch::arange(result_batch_size, torch::TensorOptions().dtype(torch::kInt64).device(device, device_id)); + auto sorted_result_psi = torch::zeros({result_batch_size, 2}, torch::TensorOptions().dtype(torch::kFloat64).device(device, device_id)); + + thrust::sort_by_key( + thrust::host, + reinterpret_cast*>(sorted_result_configs.data_ptr()), + reinterpret_cast*>(sorted_result_configs.data_ptr()) + result_batch_size, + reinterpret_cast(result_sort_index.data_ptr()), + array_less() + ); + + apply_within_kernel_interface( + /*term_number=*/term_number, + /*batch_size=*/batch_size, + /*result_batch_size=*/result_batch_size, + /*site=*/reinterpret_cast*>(site.data_ptr()), + /*kind=*/reinterpret_cast*>(kind.data_ptr()), + /*coef=*/reinterpret_cast*>(coef.data_ptr()), + /*configs=*/reinterpret_cast*>(configs.data_ptr()), + /*psi=*/reinterpret_cast*>(psi.data_ptr()), + /*result_configs=*/reinterpret_cast*>(sorted_result_configs.data_ptr()), + /*result_psi=*/reinterpret_cast*>(sorted_result_psi.data_ptr()) + ); + + auto result_psi = torch::zeros_like(sorted_result_psi); + result_psi.index_put_({result_sort_index}, sorted_result_psi); + return result_psi; +} + +template> +void add_into_heap(T* heap, std::int64_t heap_size, const T& value) { + auto compare = Compare(); + std::int64_t index = 0; + if (compare(value, heap[index])) { + } else { + if (compare(value, heap[index])) { + } else { + while (true) { + // Calculate the indices of the left and right children + std::int64_t left = (index << 1) + 1; + std::int64_t right = (index << 1) + 2; + std::int64_t left_present = left < heap_size; + std::int64_t right_present = right < heap_size; + if (left_present) { + if (right_present) { + // Both left and right children are present + if (compare(value, heap[left])) { + if (compare(value, heap[right])) { + // Both children are greater than the value, break + break; + } else { + // The left child is greater than the value, treat it as if only the right child is present + if (compare(value, heap[right])) { + break; + } else { + heap[index] = heap[right]; + index = right; + } + } + } else { + if (compare(value, heap[right])) { + // The right child is greater than the value, treat it as if only the left child is present + if (compare(value, heap[left])) { + break; + } else { + heap[index] = heap[left]; + index = left; + } + } else { + if (compare(heap[left], heap[right])) { + if (compare(value, heap[left])) { + break; + } else { + heap[index] = heap[left]; + index = left; + } + } else { + if (compare(value, heap[right])) { + break; + } else { + heap[index] = heap[right]; + index = right; + } + } + } + } + } else { + // Only the left child is present + if (compare(value, heap[left])) { + break; + } else { + if (compare(value, heap[left])) { + break; + } else { + heap[index] = heap[left]; + index = left; + } + } + } + } else { + if (right_present) { + // Only the right child is present + if (compare(value, heap[right])) { + break; + } else { + if (compare(value, heap[right])) { + break; + } else { + heap[index] = heap[right]; + index = right; + } + } + } else { + // No children are present + break; + } + } + } + heap[index] = value; + } + } +} + +template +struct array_first_double_less { + double first_double(const std::array& value) const { + double result; + for (std::int64_t i = 0; i < sizeof(double); ++i) { + reinterpret_cast(&result)[i] = reinterpret_cast(&value[0])[i]; + } + return result; + } + + bool operator()(const std::array& lhs, const std::array& rhs) const { + return first_double(lhs) < first_double(rhs); + } +}; + +template +void find_relative_kernel( + std::int64_t term_index, + std::int64_t batch_index, + std::int64_t term_number, + std::int64_t batch_size, + std::int64_t exclude_size, + const std::array* site, // term_number + const std::array* kind, // term_number + const std::array* coef, // term_number + const std::array* configs, // batch_size + const std::array* psi, // batch_size + const std::array* exclude_configs, // exclude_size + std::array* heap, + std::int64_t heap_size +) { + std::array current_configs = configs[batch_index]; + auto [success, parity] = hamiltonian_apply_kernel( + /*current_configs=*/current_configs, + /*term_index=*/term_index, + /*batch_index=*/batch_index, + /*site=*/site, + /*kind=*/kind + ); + + if (!success) { + return; + } + success = true; + std::int64_t low = 0; + std::int64_t high = exclude_size - 1; + std::int64_t mid = 0; + auto compare = array_less(); + while (low <= high) { + mid = (low + high) / 2; + if (compare(current_configs, exclude_configs[mid])) { + high = mid - 1; + } else if (compare(exclude_configs[mid], current_configs)) { + low = mid + 1; + } else { + success = false; + break; + } + } + if (!success) { + return; + } + std::int8_t sign = parity ? -1 : +1; + double real = sign * (coef[term_index][0] * psi[batch_index][0] - coef[term_index][1] * psi[batch_index][1]); + double imag = sign * (coef[term_index][0] * psi[batch_index][1] + coef[term_index][1] * psi[batch_index][0]); + // Currently, the weight is calculated as the probability of the state, but it can be changed to other values in the future. + double weight = real * real + imag * imag; + std::array value; + for (std::int64_t i = 0; i < sizeof(double) / sizeof(uint8_t); ++i) { + value[i] = reinterpret_cast(&weight)[i]; + } + for (std::int64_t i = 0; i < n_qubytes; ++i) { + value[i + sizeof(double) / sizeof(uint8_t)] = current_configs[i]; + } + add_into_heap, array_first_double_less>( + heap, + heap_size, + value + ); +} + +template +void find_relative_kernel_interface( + std::int64_t term_number, + std::int64_t batch_size, + std::int64_t exclude_size, + const std::array* site, // term_number + const std::array* kind, // term_number + const std::array* coef, // term_number + const std::array* configs, // batch_size + const std::array* psi, // batch_size + const std::array* exclude_configs, // exclude_size + std::array* heap, + std::int64_t heap_size +) { + for (std::int64_t term_index = 0; term_index < term_number; ++term_index) { + for (std::int64_t batch_index = 0; batch_index < batch_size; ++batch_index) { + find_relative_kernel( + /*term_index=*/term_index, + /*batch_index=*/batch_index, + /*term_number=*/term_number, + /*batch_size=*/batch_size, + /*exclude_size=*/exclude_size, + /*site=*/site, + /*kind=*/kind, + /*coef=*/coef, + /*configs=*/configs, + /*psi=*/psi, + /*exclude_configs=*/exclude_configs, + /*heap=*/heap, + /*heap_size=*/heap_size + ); + } + } +} + +template +auto find_relative_interface( + const torch::Tensor& configs, + const torch::Tensor& psi, + const std::int64_t count_selected, + const torch::Tensor& site, + const torch::Tensor& kind, + const torch::Tensor& coef, + const torch::Tensor& exclude_configs +) -> torch::Tensor { + std::int64_t device_id = configs.device().index(); + std::int64_t batch_size = configs.size(0); + std::int64_t term_number = site.size(0); + std::int64_t exclude_size = exclude_configs.size(0); + + TORCH_CHECK(configs.device().type() == torch::kCPU, "configs must be on CPU.") + TORCH_CHECK(configs.device().index() == device_id, "configs must be on the same device as others."); + TORCH_CHECK(configs.is_contiguous(), "configs must be contiguous.") + TORCH_CHECK(configs.dtype() == torch::kUInt8, "configs must be uint8.") + TORCH_CHECK(configs.dim() == 2, "configs must be 2D.") + TORCH_CHECK(configs.size(0) == batch_size, "configs batch size must match the provided batch_size."); + TORCH_CHECK(configs.size(1) == n_qubytes, "configs must have the same number of qubits as the provided n_qubytes."); + + TORCH_CHECK(psi.device().type() == torch::kCPU, "psi must be on CPU.") + TORCH_CHECK(psi.device().index() == device_id, "psi must be on the same device as others."); + TORCH_CHECK(psi.is_contiguous(), "psi must be contiguous.") + TORCH_CHECK(psi.dtype() == torch::kFloat64, "psi must be float64.") + TORCH_CHECK(psi.dim() == 2, "psi must be 2D.") + TORCH_CHECK(psi.size(0) == batch_size, "psi batch size must match the provided batch_size."); + TORCH_CHECK(psi.size(1) == 2, "psi must contain 2 elements for each batch."); + + TORCH_CHECK(site.device().type() == torch::kCPU, "site must be on CPU.") + TORCH_CHECK(site.device().index() == device_id, "site must be on the same device as others."); + TORCH_CHECK(site.is_contiguous(), "site must be contiguous.") + TORCH_CHECK(site.dtype() == torch::kInt16, "site must be int16.") + TORCH_CHECK(site.dim() == 2, "site must be 2D.") + TORCH_CHECK(site.size(0) == term_number, "site size must match the provided term_number."); + TORCH_CHECK(site.size(1) == max_op_number, "site must match the provided max_op_number."); + + TORCH_CHECK(kind.device().type() == torch::kCPU, "kind must be on CPU.") + TORCH_CHECK(kind.device().index() == device_id, "kind must be on the same device as others."); + TORCH_CHECK(kind.is_contiguous(), "kind must be contiguous.") + TORCH_CHECK(kind.dtype() == torch::kUInt8, "kind must be uint8.") + TORCH_CHECK(kind.dim() == 2, "kind must be 2D.") + TORCH_CHECK(kind.size(0) == term_number, "kind size must match the provided term_number."); + TORCH_CHECK(kind.size(1) == max_op_number, "kind must match the provided max_op_number."); + + TORCH_CHECK(coef.device().type() == torch::kCPU, "coef must be on CPU.") + TORCH_CHECK(coef.device().index() == device_id, "coef must be on the same device as others."); + TORCH_CHECK(coef.is_contiguous(), "coef must be contiguous.") + TORCH_CHECK(coef.dtype() == torch::kFloat64, "coef must be float64.") + TORCH_CHECK(coef.dim() == 2, "coef must be 2D.") + TORCH_CHECK(coef.size(0) == term_number, "coef size must match the provided term_number."); + TORCH_CHECK(coef.size(1) == 2, "coef must contain 2 elements for each term."); + + TORCH_CHECK(exclude_configs.device().type() == torch::kCPU, "configs must be on CPU.") + TORCH_CHECK(exclude_configs.device().index() == device_id, "configs must be on the same device as others."); + TORCH_CHECK(exclude_configs.is_contiguous(), "configs must be contiguous.") + TORCH_CHECK(exclude_configs.dtype() == torch::kUInt8, "configs must be uint8.") + TORCH_CHECK(exclude_configs.dim() == 2, "configs must be 2D.") + TORCH_CHECK(exclude_configs.size(0) == exclude_size, "configs batch size must match the provided exclude_size."); + TORCH_CHECK(exclude_configs.size(1) == n_qubytes, "configs must have the same number of qubits as the provided n_qubytes."); + + auto sorted_exclude_configs = exclude_configs.clone(torch::MemoryFormat::Contiguous); + + thrust::sort( + thrust::host, + reinterpret_cast*>(sorted_exclude_configs.data_ptr()), + reinterpret_cast*>(sorted_exclude_configs.data_ptr()) + exclude_size, + array_less() + ); + + auto result_pool = torch::zeros( + {count_selected, n_qubytes + sizeof(double) / sizeof(std::uint8_t)}, + torch::TensorOptions().dtype(torch::kUInt8).device(device, device_id) + ); + std::array* heap = + reinterpret_cast*>(result_pool.data_ptr()); + + find_relative_kernel_interface( + /*term_number=*/term_number, + /*batch_size=*/batch_size, + /*exclude_size=*/exclude_size, + /*site=*/reinterpret_cast*>(site.data_ptr()), + /*kind=*/reinterpret_cast*>(kind.data_ptr()), + /*coef=*/reinterpret_cast*>(coef.data_ptr()), + /*configs=*/reinterpret_cast*>(configs.data_ptr()), + /*psi=*/reinterpret_cast*>(psi.data_ptr()), + /*exclude_configs=*/reinterpret_cast*>(sorted_exclude_configs.data_ptr()), + /*heap=*/heap, + /*heap_size=*/count_selected + ); + + // Here, the bytes before sizeof(double) / sizeof(std::uint8_t) in result_pool are weights, and the bytes after are configs + // We need to remove items with weight 0, then sort and deduplicate the configs + + auto result_config = + result_pool.index({torch::indexing::Slice(), torch::indexing::Slice(sizeof(double) / sizeof(std::uint8_t), torch::indexing::None)}); + auto result_weight = + result_pool.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, sizeof(double) / sizeof(std::uint8_t))}); + auto nonzero = torch::any(result_weight != 0, /*dim=*/1); + auto nonzero_result_config = result_config.index({nonzero}); + auto unique_nonzero_result_config = + std::get<0>(torch::unique_dim(/*self=*/nonzero_result_config, /*dim=*/0, /*sorted=*/true, /*return_inverse=*/false, /*return_counts=*/false)); + return unique_nonzero_result_config; +} + +#ifndef N_QUBYTES +#define N_QUBYTES 0 +#endif +#ifndef PARTICLE_CUT +#define PARTICLE_CUT 0 +#endif + +#if N_QUBYTES != 0 +#define QMB_LIBRARY_HELPER(x, y) qmb_hamiltonian_##x##_##y +#define QMB_LIBRARY(x, y) QMB_LIBRARY_HELPER(x, y) +TORCH_LIBRARY_IMPL(QMB_LIBRARY(N_QUBYTES, PARTICLE_CUT), CPU, m) { + m.impl("apply_within", apply_within_interface); + m.impl("find_relative", find_relative_interface); +} +#undef QMB_LIBRARY +#undef QMB_LIBRARY_HELPER +#endif + +} // namespace qmb_hamiltonian_cpu diff --git a/qmb/hamiltonian.py b/qmb/hamiltonian.py index 8c8d378..7517099 100644 --- a/qmb/hamiltonian.py +++ b/qmb/hamiltonian.py @@ -27,6 +27,7 @@ def _load_module(cls, n_qubytes: int = 0, particle_cut: int = 0) -> object: name=name, sources=[ f"{folder}/_hamiltonian.cpp", + f"{folder}/_hamiltonian_cpu.cpp", f"{folder}/_hamiltonian_cuda.cu", ], is_python_module=n_qubytes == 0,