尝试修改hamiltonian.py，加入选择编译的功能。

qmb/hamiltonian.py

@@ -9,36 +9,76 @@
 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}"],
                 extra_cuda_cflags=["-O3", "--use_fast_math", f"-DN_QUBYTES={n_qubytes}", f"-DPARTICLE_CUT={particle_cut}"],
                 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