[Inductor] Add back the revert part (#160054)

andyanwang · pytorchmergebot · commit 05c19d1acecc · 2025-08-10T19:20:30.000Z
Add back the reverted code(#159809) as we've figured out the actual root cause of the internal test failures. Mote details in the internal diff. Rollback Plan: Differential Revision: D79776691 Pull Request resolved: #160054 Approved by: https://github.com/blaine-rister
diff --git a/torch/_dynamo/device_interface.py b/torch/_dynamo/device_interface.py
@@ -590,6 +590,10 @@ def init_device_reg() -> None:
     for i in range(torch.xpu.device_count()):
         register_interface_for_device(f"xpu:{i}", XpuInterface)
 
+    register_interface_for_device("mtia", MtiaInterface)
+    for i in range(torch.mtia.device_count()):
+        register_interface_for_device(f"mtia:{i}", MtiaInterface)
+
     register_interface_for_device("cpu", CpuInterface)
     register_interface_for_device("mps", MpsInterface)
 
diff --git a/torch/utils/_triton.py b/torch/utils/_triton.py
@@ -135,6 +135,7 @@ def _return_true(device_interface: Any) -> bool:
         "cuda": cuda_extra_check,
         "xpu": _return_true,
         "cpu": cpu_extra_check,
+        "mtia": _return_true,
     }
 
     def is_device_compatible_with_triton() -> bool:

Original file line number	Diff line number	Diff line change
`@@ -135,6 +135,7 @@ def _return_true(device_interface: Any) -> bool:`
`135`	`135`	`"cuda": cuda_extra_check,`
`136`	`136`	`"xpu": _return_true,`
`137`	`137`	`"cpu": cpu_extra_check,`
	`138`	`+ "mtia": _return_true,`
`138`	`139`	`}`
`139`	`140`
`140`	`141`	`def is_device_compatible_with_triton() -> bool:`