hamiltonian扩展条件编译 - #42
Merged
Merged
Conversation
hzhangxyz
reviewed
Jun 17, 2025
hzhangxyz
left a comment
hzhangxyz
left a comment
Member
There was a problem hiding this comment.
大部分没啥问题，就是这个verify device有点丑。
| import torch.utils.cpp_extension | ||
|
|
||
|
|
||
| def verify_device(**tensors: torch.Tensor) -> str: |
Member
There was a problem hiding this comment.
这个函数可以封装成 decorator。你要么封装成decorator，那样好看一些，要么别做这个检查，反正pytorch自己也会检查这个。
Member
There was a problem hiding this comment.
如果要做的话，大概是这样：
def verify(func):
def wrapped_func(self, *args, **kwargs):
# check device here
...
func(*args, **kwargs)
return wrapped_func但不建议你这么做。
There was a problem hiding this comment.
Pull Request Overview
This PR extends the Hamiltonian module’s conditional compilation logic by adding device support and updating the caching mechanism.
- Added a device_type parameter to control CPU vs CUDA operator loading.
- Adjusted build configurations and caching to incorporate device differentiation.
Comment on lines
+17
to
+21
| _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: |
There was a problem hiding this comment.
The caching mechanism is based on n_qubytes and device_type, but the module name also depends on particle_cut. Consider including particle_cut in the caching key to prevent potential conflicts when particle_cut varies.
hzhangxyz
reviewed
Jun 17, 2025
| """ | ||
| if configs_exclude is None: | ||
| configs_exclude = configs_i | ||
| device_type = configs_i.device.type |
| 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"}: |
| cls._hamiltonian_module[n_qubytes] = torch.utils.cpp_extension.load( | ||
| if n_qubytes == 0: | ||
| ext_paths = [] | ||
| elif device_type == "cpu": |
Member
There was a problem hiding this comment.
为了可读性，你可以这么干：
sources: list[str] = [f"{folder}/_hamiltonian.cpp"]
match (n_qubytes, device_type):
case 0, _:
pass
case _, "cpu":
sources.append(...)
case _, "cuda":
sources.append(...)
case _:
raise ValueError("Unsupported device")
This was referenced Jun 17, 2025
Merged
hzhangxyz
reviewed
Jun 18, 2025
| """ | ||
|
|
||
| _hamiltonian_module: dict[int, object] = {} | ||
| _hamiltonian_module: dict[tuple[int, str], object] = {} |
Member
There was a problem hiding this comment.
因为我平常碰到的particle_cut都是同一个值，所以当初这里没设置好，这里应该是 _hamiltonian_module: dict[tuple[str, int, int], object] = {} 才合适。
| 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: |
Member
There was a problem hiding this comment.
之前没注意，这么写是不行的，_hamiltonian.cpp 里有pytorch算子的申明（不是c/cpp的申明），如果分成cpu和cuda会导致重复编译。所以需要切割成三个情况， declaration/cpu/cuda。如果device type不是declaration，则先编译declaration的版本。
| 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 (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}" |
Member
There was a problem hiding this comment.
这个name的用处是python/pytorch模块名，而同一个算子是同一个模块的（废话），即使是适用于不同device的算子实现。
Member
|
@windy-pig 有空看看那个comment分支，更新一下这个PR吧。 |
hzhangxyz
force-pushed
the
dev/conditional-compilation
branch
2 times, most recently
from
June 26, 2025 04:03
6d8b069 to
e9e444a
Compare
hzhangxyz
pushed a commit
that referenced
this pull request
Jun 26, 2025
尝试修改hamiltonian.py，加入选择编译的功能。 删除verify_device，使用第一个传入tensor的设备类型。 按建议去除冗余混乱代码，优化可读性。
hzhangxyz
force-pushed
the
dev/conditional-compilation
branch
from
June 26, 2025 04:04
e9e444a to
1141a8e
Compare
Co-authored-by: Junkai Xu <michael17322021@163.com>
hzhangxyz
force-pushed
the
dev/conditional-compilation
branch
from
June 26, 2025 04:06
1141a8e to
49cb1b5
Compare
hzhangxyz
approved these changes
Jun 26, 2025
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
修改hamiltonian.py的逻辑，判断并选择编译对应扩展。
Checklist: