perf(dpa4): c++ infer path

Author: OutisLi

source/api_cc/include/DeepPot.h

@@ -307,6 +307,24 @@ class DeepPotBackend : public DeepBaseModelBackend {
     computew_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes,
                         coord, atype, box, fparam, aparam, atomic);
   }
+  /**
+   * @brief GPU-resident edge-input inference for the SeZM/DPA4 graph-form .pt2:
+   *given device edge tensors, write per-atom energy / force / virial back to
+   * the device output pointers. The PyTorch Exportable backend overrides this;
+   * every other backend inherits the throwing default. The signature is
+   * intentionally torch-free so the dispatcher stays backend-agnostic (no
+   * dynamic_cast into a PyTorch-heavy type, so ``libdeepmd_cc`` need not link
+   * PyTorch).
+   */
+  virtual void compute_edges_gpu(double* d_atom_energy,
+                                 double* d_force,
+                                 double* d_atom_virial,
+                                 const double* d_coord,
+                                 const int* d_atype,
+                                 const int* d_edge_index,
+                                 const double* d_edge_vec,
+                                 const int nloc,
+                                 const int nedge);
 };
 
 /**

source/api_cc/include/DeepPotPTExpt.h

@@ -318,7 +318,7 @@ class DeepPotPTExpt : public DeepPotBackend {
                          const int* d_edge_index,
                          const double* d_edge_vec,
                          const int nloc,
-                         const int nedge);
+                         const int nedge) override;
 
  private:
   bool inited;

source/api_cc/src/DeepPot.cc

@@ -564,25 +564,15 @@ template void DeepPot::compute_mixed_type<float>(
     const std::vector<float>& aparam,
     const std::vector<double>& charge_spin);
 
-void DeepPot::compute_edges_gpu(double* d_atom_energy,
-                                double* d_force,
-                                double* d_atom_virial,
-                                const double* d_coord,
-                                const int* d_atype,
-                                const int* d_edge_index,
-                                const double* d_edge_vec,
-                                const int nloc,
-                                const int nedge) {
-#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT
-  deepmd::DeepPotPTExpt* expt = dynamic_cast<deepmd::DeepPotPTExpt*>(dp.get());
-  if (expt == nullptr) {
-    throw deepmd::deepmd_exception(
-        "compute_edges_gpu is only available for the PyTorch Exportable (.pt2) "
-        "backend.");
-  }
-  expt->compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord,
-                          d_atype, d_edge_index, d_edge_vec, nloc, nedge);
-#else
+void DeepPotBackend::compute_edges_gpu(double* d_atom_energy,
+                                       double* d_force,
+                                       double* d_atom_virial,
+                                       const double* d_coord,
+                                       const int* d_atype,
+                                       const int* d_edge_index,
+                                       const double* d_edge_vec,
+                                       const int nloc,
+                                       const int nedge) {
   (void)d_atom_energy;
   (void)d_force;
   (void)d_atom_virial;
@@ -593,9 +583,28 @@ void DeepPot::compute_edges_gpu(double* d_atom_energy,
   (void)nloc;
   (void)nedge;
   throw deepmd::deepmd_exception(
-      "compute_edges_gpu requires the PyTorch Exportable backend, which is not "
-      "built.");
-#endif
+      "compute_edges_gpu (GPU-resident edge inference) is only supported by "
+      "the "
+      "PyTorch Exportable (.pt2) backend.");
+}
+
+void DeepPot::compute_edges_gpu(double* d_atom_energy,
+                                double* d_force,
+                                double* d_atom_virial,
+                                const double* d_coord,
+                                const int* d_atype,
+                                const int* d_edge_index,
+                                const double* d_edge_vec,
+                                const int nloc,
+                                const int nedge) {
+  // Polymorphic dispatch to the loaded backend: the PyTorch Exportable backend
+  // overrides ``compute_edges_gpu``; other backends inherit the throwing
+  // default. This replaces a ``dynamic_cast`` into the PyTorch-heavy
+  // ``DeepPotPTExpt``, which the "load backends as plugins" refactor made
+  // uncompilable in the backend-agnostic ``libdeepmd_cc`` (it does not link
+  // PyTorch), so the cast branch was always stubbed out.
+  dp->compute_edges_gpu(d_atom_energy, d_force, d_atom_virial, d_coord, d_atype,
+                        d_edge_index, d_edge_vec, nloc, nedge);
 }
 
 int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); }

source/api_cc/src/DeepPotPTExpt.cc

@@ -1780,9 +1780,10 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy,
         "compute_edges_gpu always returns the per-atom virial, but this .pt2 "
         "model was exported without it (do_atomic_virial=False).");
   }
-  if (!lower_input_is_edge_) {
+  if (!lower_input_is_edge_ && !lower_input_is_graph_) {
     throw deepmd::deepmd_exception(
-        "compute_edges_gpu requires an edge-input (SeZM/DPA4) .pt2 model.");
+        "compute_edges_gpu requires an edge-input (SeZM/DPA4) or graph-input "
+        "(DPA1/DPA2/DPA3) .pt2 model.");
   }
   translate_error([&] {
     const torch::Device device(torch::kCUDA, gpu_id);
@@ -1857,24 +1858,42 @@ void DeepPotPTExpt::compute_edges_gpu(double* d_atom_energy,
               .to(device);
     }
 
-    // === Step 4. Run the exported model on the device ===
-    std::vector<torch::Tensor> flat_outputs = run_model_edges(
-        coord_t, atype_t, edge_index, edge_vec, edge_scatter_index, edge_mask,
-        fparam_tensor, aparam_tensor, charge_spin_tensor);
-    std::map<std::string, torch::Tensor> output_map;
-    extract_outputs(output_map, flat_outputs);
+    // === Step 4. Run the exported model and read the per-atom outputs ===
+    // The two lower forms share the masked edge tensors but differ in both the
+    // input set and the output naming, so each form runs and unpacks itself;
+    // the result is always per-atom energy (nloc), force (nloc, 3) and virial
+    // (nloc, 9), copied out below.
+    at::Tensor ae, force_t, av;
+    std::map<std::string, torch::Tensor> out;
+    if (lower_input_is_graph_) {
+      // Graph (DPA1/DPA2/DPA3 NeighborGraph): single-frame node count and a
+      // flat node-major atype; the model returns the high-level per-atom
+      // quantities.
+      const at::Tensor n_node =
+          torch::full({1}, static_cast<std::int64_t>(nloc), opt_i64);
+      extract_outputs(
+          out, run_model_graph(atype_t.reshape({nloc}), n_node, edge_index,
+                               edge_vec, edge_mask, fparam_tensor,
+                               aparam_tensor, charge_spin_tensor));
+      ae = out["atom_energy"].reshape({nloc}).contiguous();
+      force_t = out["force"].reshape({nloc, 3}).contiguous();
+      av = out["atom_virial"].reshape({nloc, 9}).contiguous();
+    } else {
+      // Edge (SeZM/DPA4): coord + edge_scatter_index; the model returns the raw
+      // reduced-energy derivatives (force/virial per extended atom).
+      extract_outputs(
+          out, run_model_edges(coord_t, atype_t, edge_index, edge_vec,
+                               edge_scatter_index, edge_mask, fparam_tensor,
+                               aparam_tensor, charge_spin_tensor));
+      ae = out["energy"].reshape({nloc}).contiguous();
+      force_t =
+          out["energy_derv_r"].squeeze(-2).reshape({nloc, 3}).contiguous();
+      av = out["energy_derv_c"].squeeze(-2).reshape({nloc, 9}).contiguous();
+    }
 
-    // === Step 5. Copy outputs into caller GPU buffers (device-to-device) ===
-    // Per-atom energy: energy (1, nloc, 1) -> (nloc)
-    at::Tensor ae = output_map["energy"].reshape({nloc}).contiguous();
+    // === Step 5. Copy per-atom outputs into caller GPU buffers (D2D) ===
     torch::from_blob(d_atom_energy, {nloc}, opt_f64).copy_(ae);
-    // Force: energy_derv_r (1, nloc, 1, 3) -> (nloc, 3) row-major
-    at::Tensor force_t =
-        output_map["energy_derv_r"].squeeze(-2).reshape({nloc, 3}).contiguous();
     torch::from_blob(d_force, {nloc, 3}, opt_f64).copy_(force_t);
-    // Per-atom virial: energy_derv_c (1, nloc, 1, 9) -> (nloc, 9) row-major
-    at::Tensor av =
-        output_map["energy_derv_c"].squeeze(-2).reshape({nloc, 9}).contiguous();
     torch::from_blob(d_atom_virial, {nloc, 9}, opt_f64).copy_(av);
   });
 }