Decorate custom C++ ops w/compiler disable guard.

Summary: In order to work properly with `torch.compile()`, we need to decorate any function that calls a custom C++/CUDA extension with `torch.compiler.disable` so that it knows to insert a graph break.

Reviewed By: podgorskiy

Differential Revision: D59776177

fbshipit-source-id: d80eb43858836f8b8647d2a35b30d0b863989e94

drtk/edge_grad_estimator.py

@@ -16,6 +16,7 @@ from drtk.utils import index
 th.ops.load_library(edge_grad_ext.__file__)
 
 
+@th.compiler.disable
 def edge_grad_estimator(
     v_pix: th.Tensor,
     vi: th.Tensor,
@@ -126,7 +127,6 @@ def edge_grad_estimator_ref(
 
 class EdgeGradEstimatorFunction(th.autograd.Function):
     @staticmethod
-    @th.cuda.amp.custom_fwd(cast_inputs=th.float32)
     # pyre-fixme[14]: `forward` overrides method defined in `Function` inconsistently.
     def forward(
         ctx,
@@ -140,7 +140,6 @@ class EdgeGradEstimatorFunction(th.autograd.Function):
         return img
 
     @staticmethod
-    @th.cuda.amp.custom_bwd
     # pyre-fixme[14]: `backward` overrides method defined in `Function` inconsistently.
     def backward(ctx, grad_output: th.Tensor) -> Tuple[
         Optional[th.Tensor],

drtk/interpolate.py

@@ -10,6 +10,7 @@ from drtk import interpolate_ext
 th.ops.load_library(interpolate_ext.__file__)
 
 
+@th.compiler.disable
 def interpolate(
     vert_attributes: th.Tensor,
     vi: th.Tensor,

drtk/mipmap_grid_sample.py

@@ -13,6 +13,7 @@ from drtk import mipmap_grid_sampler_ext
 th.ops.load_library(mipmap_grid_sampler_ext.__file__)
 
 
+@th.compiler.disable
 def mipmap_grid_sample(
     input: List[th.Tensor],
     grid: th.Tensor,

drtk/msi.py

@@ -10,6 +10,7 @@ from drtk import msi_ext
 th.ops.load_library(msi_ext.__file__)
 
 
+@th.compiler.disable
 def msi(
     ray_o: th.Tensor,
     ray_d: th.Tensor,

drtk/rasterize.py

@@ -13,6 +13,7 @@ from drtk import rasterize_ext
 th.ops.load_library(rasterize_ext.__file__)
 
 
+@th.compiler.disable
 def rasterize(
     v: th.Tensor,
     vi: th.Tensor,
@@ -59,6 +60,7 @@ def rasterize(
     return index_img
 
 
+@th.compiler.disable
 def rasterize_with_depth(
     v: th.Tensor,
     vi: th.Tensor,

drtk/render.py

@@ -14,6 +14,7 @@ from drtk import render_ext
 th.ops.load_library(render_ext.__file__)
 
 
+@th.compiler.disable
 def render(
     v: th.Tensor,
     vi: th.Tensor,