From e4c88ea972d74e7a531dfebd2a8478dc3a6ee005 Mon Sep 17 00:00:00 2001 From: windy-pig Date: Tue, 17 Jun 2025 13:31:28 +0800 Subject: [PATCH] =?UTF-8?q?=E5=B0=9D=E8=AF=95=E4=BF=AE=E6=94=B9hamiltonian?= =?UTF-8?q?.py=EF=BC=8C=E5=8A=A0=E5=85=A5=E9=80=89=E6=8B=A9=E7=BC=96?= =?UTF-8?q?=E8=AF=91=E7=9A=84=E5=8A=9F=E8=83=BD=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- qmb/hamiltonian.py | 67 +++++++++++++++++++++++++++++++++++++--------- 1 file changed, 55 insertions(+), 12 deletions(-) diff --git a/qmb/hamiltonian.py b/qmb/hamiltonian.py index 7517099..b02f914 100644 --- a/qmb/hamiltonian.py +++ b/qmb/hamiltonian.py @@ -9,26 +9,66 @@ import torch.utils.cpp_extension +def verify_device(**tensors: torch.Tensor) -> str: + """Verify the device of tensors. + + Parameters + ---------- + **tensors : dict + Tensors passed in as argument in the main function. **At least one** tensor must be passed in. + + Returns + ------- + Literal["cuda", "cpu"] + The device type. (If all tensors are on same device type.) + + Raises + ------ + RuntimeError + When tensors are not all on the same device (CPU or CUDA). + ValueError + When no tensors are passed in. + + If all tensors are on cpu or cuda, the function returns that device string. + If they are on different devices or at least one is not on cuda or cpu, Runtime""" + if len(tensors) == 0: + raise ValueError("Must provide at least one tensor to verify") + types = set((tensor.device.type for tensor in tensors.values())) + if types - {"cuda", "cpu"}: + raise RuntimeError("Unsupported device type") + if len(types) == 2: + type_str = ", ".join((f"{name} on {str(tensor.device)}" for name, tensor in tensors.items())) + raise RuntimeError(f"Tensor device mismatch: {type_str}") + return types.pop() + + class Hamiltonian: """ The Hamiltonian type, which stores the Hamiltonian and processes iteration over each term in the Hamiltonian for given configurations. """ - _hamiltonian_module: dict[int, object] = {} + _hamiltonian_module: dict[tuple[int, str], object] = {} @classmethod - def _load_module(cls, n_qubytes: int = 0, particle_cut: int = 0) -> object: - if n_qubytes not in cls._hamiltonian_module: - name = "qmb_hamiltonian" if n_qubytes == 0 else f"qmb_hamiltonian_{n_qubytes}_{particle_cut}" + def _load_module(cls, device_type: str = "cpu", n_qubytes: int = 0, particle_cut: int = 0) -> object: + if device_type not in {"cuda", "cpu"}: + raise RuntimeError("Unsupported device type") + if (n_qubytes, device_type) not in cls._hamiltonian_module: + name = "qmb_hamiltonian" if n_qubytes == 0 else f"qmb_hamiltonian_{n_qubytes}_{particle_cut}_{device_type}" build_directory = platformdirs.user_cache_path("qmb", "kclab") / name build_directory.mkdir(parents=True, exist_ok=True) folder = os.path.dirname(__file__) - cls._hamiltonian_module[n_qubytes] = torch.utils.cpp_extension.load( + if n_qubytes == 0: + ext_paths = [] + elif device_type == "cpu": + ext_paths = [f"{folder}/_hamiltonian_cpu.cpp"] + else: + ext_paths = [f"{folder}/_hamiltonian_cuda.cu"] + cls._hamiltonian_module[(n_qubytes, device_type)] = torch.utils.cpp_extension.load( name=name, sources=[ f"{folder}/_hamiltonian.cpp", - f"{folder}/_hamiltonian_cpu.cpp", - f"{folder}/_hamiltonian_cuda.cu", + *ext_paths, ], is_python_module=n_qubytes == 0, extra_cflags=["-O3", "-ffast-math", "-march=native", f"-DN_QUBYTES={n_qubytes}", f"-DPARTICLE_CUT={particle_cut}"], @@ -36,9 +76,9 @@ def _load_module(cls, n_qubytes: int = 0, particle_cut: int = 0) -> object: build_directory=build_directory, ) if n_qubytes == 0: # pylint: disable=no-else-return - return cls._hamiltonian_module[n_qubytes] + return cls._hamiltonian_module[(n_qubytes, device_type)] else: - return getattr(torch.ops, f"qmb_hamiltonian_{n_qubytes}_{particle_cut}") + return getattr(torch.ops, f"qmb_hamiltonian_{n_qubytes}_{particle_cut}_{device_type}") @classmethod def _prepare(cls, hamiltonian: dict[tuple[tuple[int, int], ...], complex]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -96,9 +136,10 @@ def apply_within( torch.Tensor A tensor of shape [batch_size_j] representing the output amplitudes on the given configurations. """ + device_type = verify_device(configs_i=configs_i, psi_i=psi_i, configs_j=configs_j) self._prepare_data(configs_i.device) _apply_within: typing.Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] - _apply_within = getattr(self._load_module(configs_i.size(1), self.particle_cut), "apply_within") + _apply_within = getattr(self._load_module(device_type, configs_i.size(1), self.particle_cut), "apply_within") psi_j = torch.view_as_complex(_apply_within(configs_i, torch.view_as_real(psi_i), configs_j, self.site, self.kind, self.coef)) return psi_j @@ -131,9 +172,10 @@ def find_relative( """ if configs_exclude is None: configs_exclude = configs_i + device_type = verify_device(configs_i=configs_i, psi_i=psi_i, configs_exclude=configs_exclude) self._prepare_data(configs_i.device) _find_relative: typing.Callable[[torch.Tensor, torch.Tensor, int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] - _find_relative = getattr(self._load_module(configs_i.size(1), self.particle_cut), "find_relative") + _find_relative = getattr(self._load_module(device_type, configs_i.size(1), self.particle_cut), "find_relative") configs_j = _find_relative(configs_i, torch.view_as_real(psi_i), count_selected, self.site, self.kind, self.coef, configs_exclude) return configs_j @@ -151,8 +193,9 @@ def single_relative(self, configs: torch.Tensor) -> torch.Tensor: torch.Tensor A uint8 tensor of shape [batch_size, n_qubytes] representing the resulting configurations. """ + device_type = verify_device(configs=configs) self._prepare_data(configs.device) _single_relative: typing.Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] - _single_relative = getattr(self._load_module(configs.size(1), self.particle_cut), "single_relative") + _single_relative = getattr(self._load_module(device_type, configs.size(1), self.particle_cut), "single_relative") configs_result = _single_relative(configs, self.site, self.kind, self.coef) return configs_result