add 'get_active_torch_device'

anmyachev · anmyachev · commit 5df82a07db26 · 2024-12-07T19:57:16.000+01:00
Signed-off-by: Anatoly Myachev &lt;anatoly.myachev@intel.com&gt;
diff --git a/benchmarks/triton_kernels_benchmark/benchmark_driver.py b/benchmarks/triton_kernels_benchmark/benchmark_driver.py
@@ -499,6 +499,9 @@ def get_current_target(self):
         warp_size = 32
         return GPUTarget("xpu", dev_property, warp_size)
 
+    def get_active_torch_device(self):
+        return torch.device("xpu", self.get_current_device())
+
     @staticmethod
     def is_active():
         return torch.xpu.is_available()
diff --git a/python/triton/backends/driver.py b/python/triton/backends/driver.py
@@ -1,4 +1,4 @@
-from abc import ABCMeta, abstractmethod, abstractclassmethod
+from abc import ABCMeta, abstractmethod
 from typing import Callable, List, Protocol, Sequence
 
 
@@ -10,14 +10,19 @@ def __call__(self, kernel_call: Callable, *, quantiles: List[float], **kwargs) -
 
 class DriverBase(metaclass=ABCMeta):
 
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def is_active(self):
         pass
 
     @abstractmethod
     def get_current_target(self):
         pass
 
+    @abstractmethod
+    def get_active_torch_device(self):
+        pass
+
     @abstractmethod
     def get_benchmarker(self) -> Benchmarker:
         """
diff --git a/python/tutorials/01-vector-add.py b/python/tutorials/01-vector-add.py
@@ -23,8 +23,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 @triton.jit
diff --git a/python/tutorials/02-fused-softmax.py b/python/tutorials/02-fused-softmax.py
@@ -27,8 +27,7 @@
 import triton.language as tl
 from triton.runtime import driver
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 def is_hip():
@@ -122,7 +121,6 @@ def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n
 WARP_SIZE = properties["sub_group_sizes"][-1]
 WG_SIZE = properties["max_work_group_size"]
 max_num_warps = WG_SIZE // WARP_SIZE
-target = triton.runtime.driver.active.get_current_target()
 warps_per_sm = WARPS_PER_EU * EU_PER_SM
 max_num_resident_warps = NUM_SM * warps_per_sm
 kernels = {}
diff --git a/python/tutorials/03-matrix-multiplication.py b/python/tutorials/03-matrix-multiplication.py
@@ -154,8 +154,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 def is_cuda():
diff --git a/python/tutorials/04-low-memory-dropout.py b/python/tutorials/04-low-memory-dropout.py
@@ -38,8 +38,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 @triton.jit
diff --git a/python/tutorials/05-layer-norm.py b/python/tutorials/05-layer-norm.py
@@ -42,8 +42,7 @@
 except ModuleNotFoundError:
     HAS_APEX = False
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 @triton.jit
diff --git a/python/tutorials/06-fused-attention.py b/python/tutorials/06-fused-attention.py
@@ -19,8 +19,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 def is_hip():
diff --git a/python/tutorials/07-extern-functions.py b/python/tutorials/07-extern-functions.py
@@ -25,8 +25,7 @@
 
 from pathlib import Path
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 @triton.jit
diff --git a/python/tutorials/08-grouped-gemm.py b/python/tutorials/08-grouped-gemm.py
@@ -31,8 +31,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 def is_cuda():
diff --git a/python/tutorials/10-experimental-block-pointer.py b/python/tutorials/10-experimental-block-pointer.py
@@ -95,8 +95,7 @@
 import triton
 import triton.language as tl
 
-DEVICE = torch.device(triton.runtime.driver.active.get_current_target().backend,
-                      triton.runtime.driver.active.get_current_device())
+DEVICE = triton.runtime.driver.active.get_active_torch_device()
 
 
 @triton.autotune(
diff --git a/third_party/amd/backend/driver.py b/third_party/amd/backend/driver.py
@@ -502,6 +502,11 @@ def get_current_target(self):
         warp_size = device_properties['warpSize']
         return GPUTarget("hip", arch.split(':')[0], warp_size)
 
+    def get_active_torch_device(self):
+        import torch
+        # when using hip devices, the device string in pytorch is "cuda"
+        return torch.device("cuda", self.get_current_device())
+
     def get_benchmarker(self):
         from triton.testing import do_bench
         return do_bench
diff --git a/third_party/intel/backend/driver.py b/third_party/intel/backend/driver.py
@@ -558,6 +558,10 @@ def get_current_target(self):
         warp_size = 32
         return GPUTarget("xpu", dev_property, warp_size)
 
+    def get_active_torch_device(self):
+        import torch
+        return torch.device("xpu", self.get_current_device())
+
     def get_device_interface(self):
         import torch
         return torch.xpu
diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py
@@ -504,6 +504,10 @@ def get_current_target(self):
         warp_size = 32
         return GPUTarget("cuda", capability, warp_size)
 
+    def get_active_torch_device(self):
+        import torch
+        return torch.device("cuda", self.get_current_device())
+
     def get_device_interface(self):
         import torch
         return torch.cuda
