Fix PointwiseCompiler on CUDA (#203)

Author: bertmaher

functorch/_src/operator_authoring.py

@@ -316,7 +316,7 @@ def compute_code(self):
         loopnest = _te.LoopNest(_te.Block([out]), output_bufs)
 
         if self.device == "cuda" and loops:
-            flattened = _te.LoopNest.flatten(loops)
+            flattened = loopnest.flatten(loops)
             assert flattened
             inner = _te.LoopNest.split_with_mask(flattened, 512)
             assert inner

test/test_operator_authoring.py

@@ -11,6 +11,7 @@
 import unittest
 
 LLVM_ENABLED = torch._C._llvm_enabled()
+HAS_CUDA = torch.cuda.is_available()
 HAS_SYMPY = False
 try:
     import sympy
@@ -44,8 +45,8 @@ def __torch_function__(self, func, types, args=(), kwargs=None):
         return torch.zeros_like(args[0])
 
 
-class TestOperatorAuthoringCPU(JitTestCase):
-    device = "cpu"
+class TestOperatorAuthoring(JitTestCase):
+    device = None
 
     def rand(self, *args, dtype=torch.float32, **kwargs):
         return torch.randint(0, 100, args, dtype=dtype, device=self.device, **kwargs)
@@ -126,16 +127,15 @@ def example(x):
         torch.testing.assert_allclose(x + 3, graph(x))
 
 
-class TestOperatorAuthoringGPU(TestOperatorAuthoringCPU):
+@unittest.skipIf(not HAS_CUDA, "GPU tests require CUDA")
+class TestOperatorAuthoringGPU(TestOperatorAuthoring):
     device = "cuda"
 
 
-if not LLVM_ENABLED:
-    TestOperatorAuthoringCPU = None  # noqa: F811
+@unittest.skipIf(not LLVM_ENABLED, "CPU tests require LLVM")
+class TestOperatorAuthoringCPU(TestOperatorAuthoring):
+    device = "cpu"
 
-# TODO: TestOperatorAuthoringGPU is disabled because it fails on CUDAs.
-# if not torch.cuda.is_available():
-TestOperatorAuthoringGPU = None  # noqa: F811
 
 if __name__ == "__main__":
     run_tests()