Advice on clean C++ Interface involving a class with many members

LukasKoestler: <p>Dear PyTorch Team,</p>
<p>thank you very much for the library, which is very pleasant to use and has great documentation!</p>
<p>I have already written multiple CUDA extensions for PyTorch and everything works well if the C++ function takes torch::Tensors or std::vector of such as input. However, I am now in the situation where I have an <code>Index</code> class that has ca. 20 members which are each a pointer to some (GPU) memory that is owned by a <code>torch::Tensor</code>. The question is how to pass this to C++ in the best way. I have a list of options and think that the last one is best, but I would be happy about ideas and opinions.</p>
<ol>
<li>
<p>Just pass all the <code>torch::Tensor</code> as arguments or a <code>std::vector&lt;torch::Tensor&gt;</code>.<br>
1.1 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_plus_sign.png?v=12" title=":heavy_plus_sign:" class="emoji" alt=":heavy_plus_sign:" loading="lazy" width="20" height="20"> Easy implementation, easy to understand.<br>
1.2 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_minus_sign.png?v=12" title=":heavy_minus_sign:" class="emoji" alt=":heavy_minus_sign:" loading="lazy" width="20" height="20"> Not flexible. If the indexing changes slightly this change has to be manually done for all functions.</p>
</li>
<li>
<p>Implement the <code>Index</code> class in C++ and expose all necessary functionality through pybind11.<br>
2.1 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_plus_sign.png?v=12" title=":heavy_plus_sign:" class="emoji" alt=":heavy_plus_sign:" loading="lazy" width="20" height="20"> Full control in C++. Encapsulated.<br>
2.2 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_minus_sign.png?v=12" title=":heavy_minus_sign:" class="emoji" alt=":heavy_minus_sign:" loading="lazy" width="20" height="20"> No control in Python and thus no “hackability”. I don’t know if in this case memory management, reference counting, etc. works out of the box.</p>
</li>
<li>
<p><strong>Currently preferred</strong>. Implement <code>Index</code> in python (and in C++ but here with pointers only) and construct a new C++ <code>Index</code> each time a function/kernel is called.<br>
3.1 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_plus_sign.png?v=12" title=":heavy_plus_sign:" class="emoji" alt=":heavy_plus_sign:" loading="lazy" width="20" height="20"> Easy implementation, easy to understand. Can implement most things (e.g. generating the Index from data) in Python.<br>
3.2 <img src="https://discuss.pytorch.org/images/emoji/apple/heavy_minus_sign.png?v=12" title=":heavy_minus_sign:" class="emoji" alt=":heavy_minus_sign:" loading="lazy" width="20" height="20"> Maybe (very) slight runtime overhead? Have to manually make sure that pointers in C++ are valid. Other issues I forgot.</p>
</li>
</ol>
<p>Below is the test code I wrote for the third option. I used only a small <code>Index</code> but how to extend it to more fields should be clear. I would be happy to hear any suggestions and please also tell me if this is solved/documented somewhere and I missed it.</p>
<p>Thank you<br>
Lukas</p>
<h3><a name="cpp-code-1" class="anchor" href="#cpp-code-1"></a>Cpp Code</h3>
<pre data-code-wrap="cpp"><code class="lang-cpp">#include &lt;torch/extension.h&gt;


// Just POD
class Index {
public:
  int64_t n = 0;
  int64_t const *start = nullptr;
  int64_t const *size = nullptr;

  void print() const { std::cout &lt;&lt; "Index: n=" &lt;&lt; n &lt;&lt; ", start=" &lt;&lt; start &lt;&lt; ", size=" &lt;&lt; size &lt;&lt; std::endl; }
};

auto dict_from_file(std::string const &amp;fname) {
  std::ifstream input(fname, std::ios::binary);
  std::vector&lt;char&gt; bytes{std::istreambuf_iterator&lt;char&gt;(input), std::istreambuf_iterator&lt;char&gt;()};
  input.close();

  return torch::pickle_load(bytes).toGenericDict();
}

Index index_from_dict(std::unordered_map&lt;std::string, torch::Tensor&gt; const &amp;dict) {
  Index index;
  index.n = dict.at("n").item().toLong();
  index.start = dict.at("start").data_ptr&lt;int64_t&gt;();
  index.size = dict.at("size").data_ptr&lt;int64_t&gt;();
  return index;
}

Index index_from_dict(c10::Dict&lt;c10::IValue, c10::IValue&gt; const &amp;dict_in) {
  std::unordered_map&lt;std::string, torch::Tensor&gt; dict;
  for (auto const &amp;elem : dict_in) dict[elem.key().toStringRef()] = elem.value().toTensor();
  return index_from_dict(dict);
}

void test_from_python(std::unordered_map&lt;std::string, torch::Tensor&gt; dict) {
  auto index = index_from_dict(dict);
  index.print();
}

void test_from_file(std::string fname) {
  auto dict = dict_from_file(fname); // Owns the dict
  auto index = index_from_dict(dict);
  index.print();
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("test_from_python", &amp;test_from_python, "Test passing index from python");
  m.def("test_from_file", &amp;test_from_file, "Test loading index from file");
}
</code></pre>
<h3><a name="python-code-and-output-2" class="anchor" href="#python-code-and-output-2"></a>Python Code and Output</h3>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/3/83cd08af7410d8c076e72d02e83119f3376815e3.png" data-download-href="https://discuss.pytorch.org/uploads/default/83cd08af7410d8c076e72d02e83119f3376815e3" title="Screenshot from 2022-04-13 13-41-01"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83cd08af7410d8c076e72d02e83119f3376815e3_2_382x500.png" alt="Screenshot from 2022-04-13 13-41-01" data-base62-sha1="iNXPzpiMMX5Mt8F351KXScUqGxd" width="382" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83cd08af7410d8c076e72d02e83119f3376815e3_2_382x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83cd08af7410d8c076e72d02e83119f3376815e3_2_573x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83cd08af7410d8c076e72d02e83119f3376815e3_2_764x1000.png 2x" data-dominant-color="EEEFF0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot from 2022-04-13 13-41-01</span><span class="informations">882×1154 144 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>