TorchDynamo and C++ extensions

[RaulPPelaez]

We have a torch C++ Autograd extension (called get_neighbor_pairs) which provides a CPU and CUDA backends.

When trying to torch.compile a model whose forward function calls this extension:

from torchmdnet.models.utils import OptimizedDistance
import torch

model = OptimizedDistance().to("cuda")
x = torch.ones((10,3)).to("cuda")
model = torch.compile(model, backend="inductor", mode="reduce-overhead")

for _ in range(10):
    _,_,_ = model(x)

I get the following error:

Click me

Traceback (most recent call last):                                                                                                                                                                                                                                                                                                                                                                                                                                                          
  File "/shared/raul/torchmd-net/benchmarks/comp.py", line 9, in <module>                                                                                                                                                                                                                                                                                                                                                                                                                   
    _,_,_ = model(x)                                                                                                                                                                                                                                                                                                                                                                                                                                                                        
            ^^^^^^^^                                                                                                                                                                                                                                                                                                                                                                                                                                                                        
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl                                                                                                                                                                                                                                                                                                                                               
    return self._call_impl(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                                                                                                                                 
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                                                                                                                                                                                                                                                                                                 
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl                                                                                                                                                                                                                                                                                                                                                       
    return forward_call(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                                                                                                                                    
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                                                                                                                                                                                                                                                                                                                                                                                    
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 295, in _fn                                                                                                                                                                                                                                                                                                                                                              
    return fn(*args, **kwargs)                                                                                                                                                                                                                                                                                                                                                                                                                                                              
           ^^^^^^^^^^^^^^^^^^^                                                                                                                                                                                                                                                                                                                                                                                                                                                              
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1502, in _wrapped_call_impl                                                                                                                                                                                                                                                                                                                                               
    return self._call_impl(*args, **kwargs)                                                                            
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                            
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _call_impl                                                                                                         
    return forward_call(*args, **kwargs)                                                                               
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                               
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 448, in catch_errors                                                                                                       
    return callback(frame, cache_size, hooks, frame_state)                                                             
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                             
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 527, in _convert_frame                                                                                                  
    result = inner_convert(frame, cache_size, hooks, frame_state)                                                      
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                      
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 127, in _fn                                                                                                             
    return fn(*args, **kwargs)                                                                                         
           ^^^^^^^^^^^^^^^^^^^                                                                                         
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 360, in _convert_frame_assert                                                                                           
    return _compile(                                                                                                   
           ^^^^^^^^^                                                                                                   
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 180, in time_wrapper                                                                                                            
    r = func(*args, **kwargs)                                                                                          
        ^^^^^^^^^^^^^^^^^^^^^                                                                                          
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 430, in _compile                                                                                                        
    out_code = transform_code_object(code, transform)                                                                  
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                  
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1000, in transform_code_object                                                                                
    transformations(instructions, code_options)                                                                        
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 415, in transform                                                                                                       
    tracer.run()                                                                                                       
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2025, in run                                                                                                         
    super().run()                                                                                                      
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 708, in run                                                                                                          
    and self.step()                                                                                                    
        ^^^^^^^^^^^                                                                                                    
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 668, in step                                                                                                         
    getattr(self, inst.opname)(inst)                                                                                   
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 390, in wrapper                                                                                                      
    return inner_fn(self, inst)                                                                                        
           ^^^^^^^^^^^^^^^^^^^^                                                                                        
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 1697, in CALL                                                                                                        
    self.call_function(fn, args, kwargs)                                                                               
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 559, in call_function                                                                                                
    self.push(fn.call_function(self, args, kwargs))                                                                    
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                     
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/variables/torch.py", line 610, in call_function                                                                                                 
    tensor_variable = wrap_fx_proxy(                                                                                   
                      ^^^^^^^^^^^^^^                                                                                   
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1063, in wrap_fx_proxy                                                                                              
    return wrap_fx_proxy_cls(                                                                                          
           ^^^^^^^^^^^^^^^^^^                                                                                          
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/variables/builder.py", line 1099, in wrap_fx_proxy_cls                                                                                          
    example_value = get_fake_value(proxy.node, tx)                                                                     
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                     
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1298, in get_fake_value                                                                                                         
    raise TorchRuntimeError(str(e)).with_traceback(e.__traceback__) from None                                          
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1266, in get_fake_value                                                                                                         
    return wrap_fake_exception(                                                                                        
           ^^^^^^^^^^^^^^^^^^^^                                                                                        
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 860, in wrap_fake_exception                                                                                                     
    return fn()                                                                                                        
           ^^^^                                                                                                        
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1267, in <lambda>                                                                                                               
    lambda: run_node(tx.output, node, args, kwargs, nnmodule)                                                          
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                          
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1332, in run_node                                                                                                               
    raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e                                         
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 1319, in run_node                                                                                                               
    return node.target(*args, **kwargs)                                                                                
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                                
  File "/shared/raul/mambaforge/envs/torch-nightly/lib/python3.11/site-packages/torch/_ops.py", line 649, in __call__                                                                                                                         
    return self._op(*args, **kwargs or {})                                                                             
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                                                             
torch._dynamo.exc.TorchRuntimeError: Failed running call_function torchmdnet_neighbors.get_neighbor_pairs(*(), **{'strategy': 'brute', 'positions': FakeTensor(..., device='cuda:0', size=(10, 3)), 'batch': FakeTensor(..., device='cuda:0', size=(10,), dtype=torch.int64), 'max_num_pairs': 320, 'cutoff_lower': 0.0, 'cutoff_upper': 5.0, 'loop': False, 'include_transpose': True, 'box_vectors': FakeTensor(..., size=(0, 0)), 'use_periodic': False}):
The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.                                    

from user code:                                                                                                        
   File "/shared/raul/torchmd-net/torchmdnet/models/utils.py", line 240, in forward                                    
    edge_index, edge_vec, edge_weight, num_pairs = get_neighbor_pairs_kernel(                                          

Being, to me, the most significant part:

torch._dynamo.exc.TorchRuntimeError: Failed running call_function torchmdnet_neighbors.get_neighbor_pairs(*(), **{'strategy': 'brute', 'positions': FakeTensor(..., device='cuda:0', size=(10, 3)), 'batch': FakeTensor(..., device='cuda:0', size=(10,), dtype=torch.int64), 'max_num_pairs': 320, 'cutoff_lower': 0.0, 'cutoff_upper': 5.0, 'loop': False, 'include_transpose': True, 'box_vectors': FakeTensor(..., size=(0, 0)), 'use_periodic': False}):
The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.                                    

from user code:                                                                                                        
   File "/shared/raul/torchmd-net/torchmdnet/models/utils.py", line 240, in forward                                    
    edge_index, edge_vec, edge_weight, num_pairs = get_neighbor_pairs_kernel(          

The extension is defined as:

torchmd-net/torchmdnet/neighbors/neighbors.cpp

Lines 1 to 5 in a116847

	#include <torch/extension.h>
	TORCH_LIBRARY(torchmdnet_neighbors, m) {
	m.def("get_neighbor_pairs(str strategy, Tensor positions, Tensor batch, Tensor box_vectors, bool use_periodic, Scalar cutoff_lower, Scalar cutoff_upper, Scalar max_num_pairs, bool loop, bool include_transpose) -> (Tensor neighbors, Tensor distances, Tensor distance_vecs, Tensor num_pairs)");
	}

And the CUDA implementation is here

torchmd-net/torchmdnet/neighbors/neighbors_cuda.cu

Lines 74 to 89 in a116847

	TORCH_LIBRARY_IMPL(torchmdnet_neighbors, AutogradCUDA, m) {
	m.impl("get_neighbor_pairs",
	[](const std::string& strategy, const Tensor& positions, const Tensor& batch,
	const Tensor& box_vectors, bool use_periodic, const Scalar& cutoff_lower,
	const Scalar& cutoff_upper, const Scalar& max_num_pairs, bool loop,
	bool include_transpose) {
	auto final_strategy = strategy;
	if (positions.size(0) >= 32768 && strategy == "brute") {
	final_strategy = "shared";
	}
	auto result = NeighborAutograd::apply(final_strategy, positions, batch, box_vectors,
	use_periodic, cutoff_lower, cutoff_upper,
	max_num_pairs, loop, include_transpose);
	return std::make_tuple(result[0], result[1], result[2], result[3]);
	});
	}

I jit compile this function when import happens:

torchmd-net/torchmdnet/neighbors/__init__.py

Lines 2 to 16 in a116847

	import torch
	from torch.utils import cpp_extension
	def compile_extension():
	src_dir = os.path.dirname(__file__)
	sources = ["neighbors.cpp", "neighbors_cpu.cpp"] + (
	["neighbors_cuda.cu", "backwards.cu"] if torch.cuda.is_available() else []
	)
	sources = [os.path.join(src_dir, name) for name in sources]
	cpp_extension.load(
	name="torchmdnet_neighbors", sources=sources, is_python_module=False
	)
	compile_extension()
	get_neighbor_pairs_kernel = torch.ops.torchmdnet_neighbors.get_neighbor_pairs

Eventually I call this extension in the forward function of the model in the example above:

torchmd-net/torchmdnet/models/utils.py

Line 110 in a116847

class OptimizedDistance(torch.nn.Module):

Which stores the extension function as self.kernel

torchmd-net/torchmdnet/models/utils.py

Lines 234 to 237 in a116847

	edge_index, edge_vec, edge_weight, num_pairs = self.kernel(
	strategy=self.strategy,
	positions=pos,
	batch=batch,