[AOTI][reland] Update test runner to use the new APIs (pytorch#149412)

Summary: Reland pytorch#147105. Switch to the newer aoti_compile_and_package APIs. Some tests still kept using legacy APIs, and will follow up with internal test refactoring.

Differential Revision: [D71470265](https://our.internmc.facebook.com/intern/diff/D71470265)
Pull Request resolved: pytorch#149412
Approved by: https://github.com/yushangdi

Author: desertfire

test/distributed/test_c10d_functional_native.py

@@ -747,7 +747,7 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         assert "= torch.ops._c10d_functional.wait_tensor.default" not in code
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (arg,))
+        AOTIRunnerUtil.run(func, (arg,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -793,7 +793,7 @@ def func(args: list[torch.Tensor]) -> torch.Tensor:
         assert "= torch.ops._c10d_functional.wait_tensor.default" not in code
 
         # Test aoti
-        out = AOTIRunnerUtil.run("cuda", func, (args,))  # noqa: F841
+        out = AOTIRunnerUtil.run(func, (args,))  # noqa: F841
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -905,7 +905,7 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         assert "= torch.ops._c10d_functional.wait_tensor.default" not in code
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (arg,))
+        AOTIRunnerUtil.run(func, (arg,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -939,7 +939,7 @@ def func(args: list[torch.Tensor]) -> torch.Tensor:
         )
 
         # Test aoti
-        out = AOTIRunnerUtil.run("cuda", func, (args,))  # noqa: F841
+        out = AOTIRunnerUtil.run(func, (args,))  # noqa: F841
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "This is a GPU test!")
@@ -961,7 +961,7 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         )
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (arg,))
+        AOTIRunnerUtil.run(func, (arg,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -987,7 +987,7 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         )
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (arg,))
+        AOTIRunnerUtil.run(func, (arg,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -1023,7 +1023,7 @@ def func(args: list[torch.Tensor]) -> torch.Tensor:
         )
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (args,))
+        AOTIRunnerUtil.run(func, (args,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
@@ -1108,7 +1108,7 @@ def func(arg: torch.Tensor) -> torch.Tensor:
         )
 
         # Test aoti
-        AOTIRunnerUtil.run("cuda", func, (arg,))
+        AOTIRunnerUtil.run(func, (arg,))
         torch.cuda.synchronize()
 
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")