C++/CUDA pre-built binaries

tmol/database/scoring/omega_bbdep.py

@@ -39,5 +39,4 @@ def from_file(cls, fname: str):
                 OmegaBBDepTables,
             ]
         ):
-            print("safe globals: ", torch.serialization.get_safe_globals())
             return torch.load(fname)

tmol/io/details/compiled/compiled.ops.cpp

@@ -226,6 +226,12 @@ TORCH_LIBRARY_(TORCH_EXTENSION_NAME, m) {
   m.def("resolve_his_taut", &resolve_his_tautomerization);
 }
 
+// #define PYBIND11_MODULE_(ns, m) PYBIND11_MODULE(ns, m)
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("gen_pose_leaf_atoms", &pose_leaf_atom_gen_op);
+  m.def("resolve_his_taut", &resolve_his_tautomerization);
+}
+
 }  // namespace compiled
 }  // namespace details
 }  // namespace io

tmol/io/details/compiled/compiled.py

@@ -1,23 +1,19 @@
-import torch
-from tmol.utility.cpp_extension import load, modulename, relpaths, cuda_if_available
+from tmol.utility.cpp_extension import relpaths, TorchOpLoader
 
-load(
-    modulename(__name__),
-    cuda_if_available(
-        relpaths(
-            __file__,
-            [
-                "compiled.ops.cpp",
-                "gen_pose_leaf_atoms.cpu.cpp",
-                "gen_pose_leaf_atoms.cuda.cu",
-                "resolve_his_taut.cpu.cpp",
-                "resolve_his_taut.cuda.cu",
-            ],
-        )
-    ),
-    is_python_module=False,
+sources = relpaths(
+    __file__,
+    [
+        "compiled.ops.cpp",
+        "gen_pose_leaf_atoms.cpu.cpp",
+        "gen_pose_leaf_atoms.cuda.cu",
+        "resolve_his_taut.cpu.cpp",
+        "resolve_his_taut.cuda.cu",
+    ],
 )
 
-_ops = getattr(torch.ops, modulename(__name__))
-gen_pose_leaf_atoms = _ops.gen_pose_leaf_atoms
-resolve_his_taut = _ops.resolve_his_taut
+functions = ["gen_pose_leaf_atoms", "resolve_his_taut"]
+
+loader = TorchOpLoader(__name__, sources, functions)
+
+gen_pose_leaf_atoms = loader.gen_pose_leaf_atoms
+resolve_his_taut = loader.resolve_his_taut

tmol/kinematics/compiled/__init__.py

@@ -1,4 +1,3 @@
-from .compiled_ops import forward_kin_op
-from .compiled_inverse_kin import inverse_kin
+from .compiled_ops import forward_kin_op, inverse_kin
 
 __all__ = ["forward_kin_op", "inverse_kin"]

tmol/kinematics/compiled/compiled_ops.cpp

@@ -19,6 +19,66 @@ using torch::autograd::AutogradContext;
 using torch::autograd::Function;
 using torch::autograd::tensor_list;
 
+class InverseKinematicOp
+    : public torch::autograd::Function<InverseKinematicOp> {
+ public:
+  static Tensor forward(
+      AutogradContext* ctx,
+      Tensor coords,
+      Tensor parent,
+      Tensor frame_x,
+      Tensor frame_y,
+      Tensor frame_z,
+      Tensor doftype) {
+    at::Tensor ret;
+
+    using Int = int32_t;
+
+    TMOL_DISPATCH_FLOATING_DEVICE(coords.options(), "inverse_kin_op", ([&] {
+                                    using Real = scalar_t;
+                                    constexpr tmol::Device Dev = device_t;
+
+                                    auto result =
+                                        InverseKinDispatch<Dev, Real, Int>::f(
+                                            TCAST(coords),
+                                            TCAST(parent),
+                                            TCAST(frame_x),
+                                            TCAST(frame_y),
+                                            TCAST(frame_z),
+                                            TCAST(doftype));
+                                    ret = result.tensor;
+                                  }));
+
+    // ctx->save_for_backward({HTs, dofs, nodes_b, scans_b, gens_b, kintree});
+
+    return ret;
+  }
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs) {
+    // auto saved = ctx->get_saved_variables();
+
+    return {
+        torch::Tensor(),
+        torch::Tensor(),
+        torch::Tensor(),
+        torch::Tensor(),
+        torch::Tensor(),
+        torch::Tensor(),
+    };
+  }
+};
+
+Tensor inverse_kin_op(
+    Tensor coords,
+    Tensor parent,
+    Tensor frame_x,
+    Tensor frame_y,
+    Tensor frame_z,
+    Tensor doftype) {
+  return InverseKinematicOp::apply(
+      coords, parent, frame_x, frame_y, frame_z, doftype);
+}
+
 class KinematicOp : public torch::autograd::Function<KinematicOp> {
  public:
   static Tensor forward(
@@ -611,5 +671,24 @@ TORCH_LIBRARY_(TORCH_EXTENSION_NAME, m) {
   m.def("minimizer_map_from_movemap", &minimizer_map_from_movemap);
 }
 
+// #define PYBIND11_MODULE_(ns, m) PYBIND11_MODULE(ns, m)
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward_kin_op", &kinematic_op);
+  m.def("forward_only_op", &forward_only_op);
+  m.def("get_kfo_indices_for_atoms", &get_kfo_indices_for_atoms);
+  m.def("get_kfo_atom_parents", &get_kfo_atom_parents);
+  m.def("get_children", &get_children);
+  m.def("get_id_and_frame_xyz", &get_id_and_frame_xyz);
+  m.def("calculate_ff_edge_delays", &calculate_ff_edge_delays);
+  m.def("get_jump_atom_indices", &get_jump_atom_indices);
+  m.def(
+      "get_block_parent_connectivity_from_toposort",
+      &get_block_parent_connectivity_from_toposort);
+  m.def("get_kinforest_scans_from_stencils", &get_scans2);
+  m.def("get_kinforest_scans_from_stencils2", &get_scans2);
+  m.def("minimizer_map_from_movemap", &minimizer_map_from_movemap);
+  m.def("inverse_kin", &inverse_kin_op);
+}
+
 }  // namespace kinematics
 }  // namespace tmol

tmol/kinematics/compiled/compiled_ops.py

@@ -1,26 +1,39 @@
-import torch
-from tmol.utility.cpp_extension import load, relpaths, modulename, cuda_if_available
+from tmol.utility.cpp_extension import relpaths, TorchOpLoader
 
-load(
-    modulename(__name__),
-    cuda_if_available(
-        relpaths(__file__, ["compiled_ops.cpp", "compiled.cpu.cpp", "compiled.cuda.cu"])
-    ),
-    is_python_module=False,
+sources = relpaths(
+    __file__, ["compiled_ops.cpp", "compiled.cpu.cpp", "compiled.cuda.cu"]
 )
 
-_ops = getattr(torch.ops, modulename(__name__))
-forward_kin_op = _ops.forward_kin_op
-forward_only_op = _ops.forward_only_op
-get_kfo_indices_for_atoms = _ops.get_kfo_indices_for_atoms
-get_kfo_atom_parents = _ops.get_kfo_atom_parents
-get_children = _ops.get_children
-get_id_and_frame_xyz = _ops.get_id_and_frame_xyz
-calculate_ff_edge_delays = _ops.calculate_ff_edge_delays
-get_jump_atom_indices = _ops.get_jump_atom_indices
+functions = [
+    "forward_kin_op",
+    "forward_only_op",
+    "get_kfo_indices_for_atoms",
+    "get_kfo_atom_parents",
+    "get_children",
+    "get_id_and_frame_xyz",
+    "calculate_ff_edge_delays",
+    "get_jump_atom_indices",
+    "get_block_parent_connectivity_from_toposort",
+    "get_kinforest_scans_from_stencils",
+    "get_kinforest_scans_from_stencils2",
+    "minimizer_map_from_movemap",
+    "inverse_kin",
+]
+
+loader = TorchOpLoader(__name__, sources, functions)
+
+forward_kin_op = loader.forward_kin_op
+forward_only_op = loader.forward_only_op
+get_kfo_indices_for_atoms = loader.get_kfo_indices_for_atoms
+get_kfo_atom_parents = loader.get_kfo_atom_parents
+get_children = loader.get_children
+get_id_and_frame_xyz = loader.get_id_and_frame_xyz
+calculate_ff_edge_delays = loader.calculate_ff_edge_delays
+get_jump_atom_indices = loader.get_jump_atom_indices
 get_block_parent_connectivity_from_toposort = (
-    _ops.get_block_parent_connectivity_from_toposort
+    loader.get_block_parent_connectivity_from_toposort
 )
-get_kinforest_scans_from_stencils = _ops.get_kinforest_scans_from_stencils
-get_kinforest_scans_from_stencils2 = _ops.get_kinforest_scans_from_stencils2
-minimizer_map_from_movemap = _ops.minimizer_map_from_movemap
+get_kinforest_scans_from_stencils = loader.get_kinforest_scans_from_stencils
+get_kinforest_scans_from_stencils2 = loader.get_kinforest_scans_from_stencils2
+minimizer_map_from_movemap = loader.minimizer_map_from_movemap
+inverse_kin = loader.inverse_kin

tmol/numeric/bspline_compiled/compiled.py

@@ -1,11 +1,21 @@
-from tmol.utility.cpp_extension import load, relpaths, modulename
+from tmol.utility.cpp_extension import relpaths, TorchOpLoader
 
-# only CPU is exposed via pybind
-_compiled = load(modulename(__name__), relpaths(__file__, ["bspline.pybind.cpp"]))
+sources = relpaths(__file__, ["bspline.pybind.cpp"])
 
-computeCoeffs2 = _compiled.computeCoeffs2
-interpolate2 = _compiled.interpolate2
-computeCoeffs3 = _compiled.computeCoeffs3
-interpolate3 = _compiled.interpolate3
-computeCoeffs4 = _compiled.computeCoeffs4
-interpolate4 = _compiled.interpolate4
+functions = [
+    "computeCoeffs2",
+    "interpolate2",
+    "computeCoeffs3",
+    "interpolate3",
+    "computeCoeffs4",
+    "interpolate4",
+]
+
+loader = TorchOpLoader(__name__, sources, functions)
+
+computeCoeffs2 = loader.computeCoeffs2
+interpolate2 = loader.interpolate2
+computeCoeffs3 = loader.computeCoeffs3
+interpolate3 = loader.interpolate3
+computeCoeffs4 = loader.computeCoeffs4
+interpolate4 = loader.interpolate4

tmol/pack/rotamer/build_rotamers.py

@@ -414,7 +414,7 @@ def _t(arr):
 def measure_dofs_from_orig_coords(
     coords: Tensor[torch.float32][:, :, :], kinforest: KinForest
 ):
-    from tmol.kinematics.compiled.compiled_inverse_kin import inverse_kin
+    from tmol.kinematics.compiled.compiled_ops import inverse_kin
 
     kinforest_coords = coords.view(-1, 3)[kinforest.id.to(torch.int64)]
     kinforest_coords[0, :] = 0  # reset root

tmol/pack/rotamer/dunbrack/compiled.ops.cpp

@@ -150,6 +150,11 @@ TORCH_LIBRARY_(TORCH_EXTENSION_NAME, m) {
   m.def("dun_sample_chi", &dun_sample_chi<score::common::ComplexDispatch>);
 }
 
+// #define PYBIND11_MODULE_(ns, m) PYBIND11_MODULE(ns, m)
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("dun_sample_chi", &dun_sample_chi<score::common::ComplexDispatch>);
+}
+
 }  // namespace dunbrack
 }  // namespace rotamer
 }  // namespace pack

tmol/pack/rotamer/dunbrack/compiled.py

@@ -1,15 +1,11 @@
-import torch
-from tmol.utility.cpp_extension import load, relpaths, modulename, cuda_if_available
+from tmol.utility.cpp_extension import relpaths, TorchOpLoader
 
-
-load(
-    modulename(__name__),
-    cuda_if_available(
-        relpaths(__file__, ["compiled.ops.cpp", "compiled.cpu.cpp", "compiled.cuda.cu"])
-    ),
-    is_python_module=False,
+sources = relpaths(
+    __file__, ["compiled.ops.cpp", "compiled.cpu.cpp", "compiled.cuda.cu"]
 )
 
-_ops = getattr(torch.ops, modulename(__name__))
+functions = ["dun_sample_chi"]
+
+loader = TorchOpLoader(__name__, sources, functions)
 
-dun_sample_chi = _ops.dun_sample_chi
+dun_sample_chi = loader.dun_sample_chi

tmol/pack/rotamer/single_residue_kinforest.py

@@ -7,7 +7,7 @@
 
 from tmol.kinematics.scan_ordering import KinForestScanOrdering
 from tmol.kinematics.old.builder import _KinematicBuilder
-from tmol.kinematics.compiled.compiled_inverse_kin import inverse_kin
+from tmol.kinematics.compiled.compiled_ops import inverse_kin
 
 from tmol.chemical.restypes import RefinedResidueType
 from tmol.pose.packed_block_types import PackedBlockTypes

tmol/pack/sim_anneal/compiled/compiled.ops.cpp

@@ -334,6 +334,25 @@ TORCH_LIBRARY_(TORCH_EXTENSION_NAME, m) {
   m.def("run_sim_annealing", &run_sim_annealing);
 }
 
+// #define PYBIND11_MODULE_(ns, m) PYBIND11_MODULE(ns, m)
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def(
+      "pick_random_rotamers",
+      &pick_random_rotamers<tmol::score::common::ForallDispatch>);
+  m.def(
+      "metropolis_accept_reject",
+      &metropolis_accept_reject<tmol::score::common::ForallDispatch>);
+  m.def("create_sim_annealer", &create_sim_annealer);
+  m.def("delete_sim_annealer", &delete_sim_annealer);
+  m.def(
+      "register_standard_random_rotamer_picker",
+      &register_standard_random_rotamer_picker);
+  m.def(
+      "register_standard_metropolis_accept_or_rejector",
+      &register_standard_metropolis_accept_or_rejector);
+  m.def("run_sim_annealing", &run_sim_annealing);
+}
+
 }  // namespace compiled
 }  // namespace sim_anneal
 }  // namespace pack