Add torch dispatcher op for tf32_hc_prenorm_gemm

[leevan]

What

Registers deep_gemm::tf32_hc_prenorm_gemm as a torch dispatcher op (torch.ops.deep_gemm.tf32_hc_prenorm_gemm) in addition to the existing pybind binding.

Why

External C++ (and Python) callers can then invoke this kernel through the PyTorch dispatcher without including DeepGEMM headers or linking against DeepGEMM — e.g. a host-side C++ extension can resolve it via c10::Dispatcher::singleton().findSchemaOrThrow("deep_gemm::tf32_hc_prenorm_gemm", ""). The schema is compiled into the _C extension alongside the pybind bindings, so the op is available as soon as import deep_gemm loads _C.

Details

• New header csrc/apis/torch_ops.hpp:
  • A thin forwarder to deep_gemm::hyperconnection::tf32_hc_prenorm_gemm, converting std::optional<int64_t> → std::optional<int>.
  • TORCH_LIBRARY(deep_gemm) schema: tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(s!) sqr_sum, int? num_splits=None) -> () — d and sqr_sum are marked mutable (written in place), num_splits defaults to None to match the pybind signature.
  • TORCH_LIBRARY_IMPL(deep_gemm, CUDA) registering the CUDA implementation.
• csrc/python_api.cpp: include the new header so the schema is compiled into _C.

Notes

• The TORCH_LIBRARY handle is named lib (not m) so the .pyi stub generator, which scans csrc for pybind-style m.def(...), skips these dispatcher schemas.
• Only a CUDA impl is registered (no Meta/fake kernel); the op is intended for eager execution. Happy to add a Meta kernel if there is interest in tracing through it under torch.compile.
• This wraps a single API to establish the pattern; more host APIs can be exposed the same way later if useful.

[ds-review-bot]

🟡 warning: TORCH_LIBRARY / TORCH_LIBRARY_IMPL 会在每个包含该头文件的翻译单元中生成静态注册对象，#pragma once 无法跨 TU 保护。当前 setup.py 只编译 csrc/python_api.cpp 单个 TU，因此是安全的；但若将来 sources 增加第二个 .cpp 并 include 此头，会在加载时触发 'Only a single TORCH_LIBRARY can be used to register the namespace deep_gemm' 错误。建议在头文件顶部注释中明确“本头文件只能被恰好一个 TU 包含”，或将注册块移入独立的 torch_ops.cpp 并加入 setup.py 的 sources。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: 转发函数使用 std::optional / int64_t，但未显式 #include 与 ，目前依赖 torch/library.h 的传递包含。建议显式包含以增强健壮性。另外该文件使用 at::Tensor，而仓库其余 API 统一使用 torch::Tensor，可考虑保持一致（二者为同一类型，仅为风格问题）。

_{🤖 v5}

[ds-review-bot]

🔵 suggestion: TORCH_LIBRARY(deep_gemm, lib) cannot be used more than once per translation unit: it expands to a static TORCH_LIBRARY_init_deep_gemm and TORCH_LIBRARY_static_init_deep_gemm. Since the header is included from python_api.cpp and the description states more host APIs will be exposed "the same way later", a second header using TORCH_LIBRARY(deep_gemm, ...) in the same TU will fail to compile with a redefinition error. Consider using TORCH_LIBRARY_FRAGMENT(deep_gemm, lib) here (the documented mechanism for splitting schema registration across files), or centralizing all deep_gemm schemas in a single registration block.

_{🤖 v4}

[ds-review-bot]

🔵 suggestion: The schema exposes num_splits as int? (64-bit) but the forwarder narrows it with static_cast<int>. A value larger than INT_MAX would silently truncate/overflow before reaching the kernel. num_splits is realistically small, but a checked conversion (or asserting num_splits <= INT_MAX) would avoid a silent wraparound.

_{🤖 v4}

[ds-review-bot]

🤖 ds-review-bot Code Review

v6

新增 dispatcher 入口绕过了共享 JIT 状态原有的 GIL 保护，引入并发调用时的数据竞争。当前环境缺少 PyTorch 和 CUDA 编译工具，未运行构建或 GPU 测试。

v5

本 MR 新增 csrc/apis/torch_ops.hpp，将 deep_gemm::hyperconnection::tf32_hc_prenorm_gemm 通过 TORCH_LIBRARY / TORCH_LIBRARY_IMPL 注册为 torch dispatcher op（torch.ops.deep_gemm.tf32_hc_prenorm_gemm），并在 csrc/python_api.cpp 中 include 该头文件使 schema 编译进 _C。静态审查结论：实现正确、与描述一致。(1) 转发函数签名与 pybind 目标兼容：at::Tensor& 可绑定 const torch::Tensor&，int? 正确映射为 std::optional<int64_t> 并转换为 std::optional<int>；(2) Tensor(d!)/Tensor(s!) 可变标注与 kernel 就地写入语义一致，num_splits=None 与 pybind 默认值 std::nullopt 对齐；(3) TORCH_LIBRARY 放在头文件中在当前构建下是安全的——setup.py 只编译单个 TU（csrc/python_api.cpp），且仓库中不存在其他 TORCH_LIBRARY(deep_gemm, ...)，不会出现重复注册；(4) scripts/generate_pyi.py 以字面量 'm.def(' 匹配，lib.def( 不会被误采集，头文件注释中的 m.def(...) 以 // 开头也会被生成器跳过；(5) csrc/python_api.cpp 中的 include 位置合理，注册随 _C 加载完成，符合“import deep_gemm 即可用”的预期。未做编译/运行验证（环境仅允许 git 命令）。整体建议：可合并，附带若干非阻塞建议。

v4

This change registers the existing deep_gemm::hyperconnection::tf32_hc_prenorm_gemm host API as a torch dispatcher op (torch.ops.deep_gemm.tf32_hc_prenorm_gemm) so external C++/Python callers can invoke it without including DeepGEMM headers or linking against DeepGEMM. The new csrc/apis/torch_ops.hpp is a small, well-scoped forwarder that converts std::optional<int64_t> to std::optional<int>, marks the in-place outputs d/sqr_sum as mutable (Tensor(d!), Tensor(s!)), keeps num_splits optional to mirror the pybind signature, and registers a CUDA implementation; csrc/python_api.cpp includes the header so the schema lands in _C. The implementation looks correct for eager CUDA execution and the lib handle name cleanly avoids the m.def(...) based .pyi stub generator. Main follow-ups are around extensibility of the registration mechanism, missing test coverage at the dispatcher level, and the absence of a Meta/fake kernel for torch.compile tracing (acknowledged in the description).

Files reviewed: 2
Issues found: 🔴 1 critical | 🟡 1 warning | 🔵 7 suggestion
Inline comments posted: 8
General comments (无法定位到 diff): 1

---

📍 未定位到 diff 的评论

🔵 suggestion tests/test_hyperconnection.py:L30: No test exercises the newly added dispatcher entry point. test_hyperconnection.py only calls the pybind function deep_gemm.tf32_hc_prenorm_gemm, so the torch-op schema/registration could regress (e.g. wrong argument types, missing mutable annotations, or the op not resolving via c10::Dispatcher::findSchemaOrThrow) without any test noticing. Please add coverage that resolves and runs the op through the dispatcher, e.g. torch.ops.deep_gemm.tf32_hc_prenorm_gemm(a, b, d, s, num_splits=...), mirroring the existing parameterized cases. _{🤖 v4}

[leevan]

Thanks for the detailed review. Pushed a follow-up commit (9f07225) addressing most of the comments:

• 🟡 multi-TU TORCH_LIBRARY: switched to TORCH_LIBRARY_FRAGMENT(deep_gemm, lib), which is the documented mechanism for splitting schema registration across files. This directly supports exposing more deep_gemm ops the same way later without a duplicate-namespace error. Also added a header comment stating the header must be included by exactly one TU.

• 🔵 num_splits int64→int narrowing: added a TORCH_CHECK that validates the value is in [1, INT_MAX] before the static_cast<int>, so an out-of-range value fails loudly instead of silently wrapping.

• 🔵 explicit includes / style: added <cstdint>, <limits>, <optional> and switched at::Tensor → torch::Tensor to match the rest of the APIs.

• 🔵 test coverage: added test_hc_prenorm_gemm_torch_op in tests/test_hyperconnection.py, which runs the op through the dispatcher (torch.ops.deep_gemm.tf32_hc_prenorm_gemm) and compares it against the pybind path for both num_splits=None and an explicit split. This guards the schema/forwarder signature against future pybind changes.

• 🔵 Meta/fake kernel: documented the eager-only limitation in the header (tracing under torch.compile/make_fx fails with "not implemented for Meta"). Happy to add a Meta kernel in a follow-up if torch.compile support for external callers is desired.

• 🔴 GIL / shared JIT state: this is a real window — the dispatcher entry can be called without the GIL, unlike the pybind binding, so concurrent first-time compilation of new shapes depends on DeepJIT's own thread-safety. I've documented the contract in the header (callers should warm up the cache before concurrent use). I'd prefer not to unconditionally gil_scoped_acquire here, since (a) it reintroduces a Python dependency that this entry point is specifically meant to avoid, and (b) the right fix may be to make the JIT cache thread-safe rather than serialize on the GIL. Could a maintainer confirm whether DeepJIT's compile cache is expected to be thread-safe? Happy to follow whichever direction you prefer — add a GIL guard here, or track cache-level locking as a separate issue.

Update: verified locally on an SM90 GPU (CUDA 13.2, PyTorch 2.11, Python 3.12, cxx11abi).

• _C builds cleanly (only a benign visibility warning on GilScopedRelease).
• import deep_gemm loads and torch.ops.deep_gemm.tf32_hc_prenorm_gemm resolves with the expected schema deep_gemm::tf32_hc_prenorm_gemm(Tensor a, Tensor b, Tensor(d!) d, Tensor(s!) sqr_sum, int? num_splits=None) -> ().
• tests/test_hyperconnection.py passes, including the new test_hc_prenorm_gemm_torch_op: the dispatcher path matches the pybind path for num_splits = None and 16 (calc_diff < 1e-10).