Parse architecture from PyTorch instead of hard coding (#2995)

With pytorch/pytorch#138186, `architecture` is
added to XPU device property.
Instead of hard coding `pvc` when invoking `ocloc`, this PR changed to
dynamically passing the device architecture parsed.

---------

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

third_party/intel/backend/arch_parser.c

@@ -0,0 +1,57 @@
+//===- arch_parser.c ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <sycl/sycl.hpp>
+
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#include <Python.h>
+#include <numpy/arrayobject.h>
+
+static PyObject *parseDeviceArch(PyObject *self, PyObject *args) {
+  uint64_t dev_arch;
+  assert(PyArg_ParseTuple(args, "K", &dev_arch) && "Expected an integer");
+
+  sycl::ext::oneapi::experimental::architecture sycl_arch =
+      static_cast<sycl::ext::oneapi::experimental::architecture>(dev_arch);
+  // FIXME: Add support for more architectures.
+  std::string arch = "";
+  switch (sycl_arch) {
+  case sycl::ext::oneapi::experimental::architecture::intel_gpu_pvc:
+    arch = "pvc";
+    break;
+  case sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g21:
+    arch = "bmg";
+    break;
+  case sycl::ext::oneapi::experimental::architecture::intel_gpu_lnl_m:
+    arch = "lnl";
+    break;
+  default:
+    printf("sycl_arch = %d", sycl_arch);
+  }
+
+  return Py_BuildValue("s", arch.c_str());
+}
+
+static PyMethodDef ModuleMethods[] = {
+    {"parse_device_arch", parseDeviceArch, METH_VARARGS,
+     "parse device architecture"},
+    {NULL, NULL, 0, NULL} // sentinel
+};
+
+static struct PyModuleDef ModuleDef = {PyModuleDef_HEAD_INIT, "arch_utils",
+                                       NULL, // documentation
+                                       -1,   // size
+                                       ModuleMethods};
+
+PyMODINIT_FUNC PyInit_arch_utils(void) {
+  if (PyObject *m = PyModule_Create(&ModuleDef)) {
+    PyModule_AddFunctions(m, ModuleMethods);
+    return m;
+  }
+  return NULL;
+}

third_party/intel/backend/compiler.py

@@ -1,5 +1,6 @@
 from triton.backends.compiler import BaseBackend
 from triton._C.libtriton import ir, passes, llvm, intel
+from triton.backends.intel.driver import compile_module_from_src
 
 from dataclasses import dataclass
 import functools
@@ -96,6 +97,7 @@ def get_ops_per_channel(lhs_type, rhs_type):
 
 
 class XPUBackend(BaseBackend):
+    device_props: dict = {}
 
     # AdvancedPath pass pipeline for kernels using block pointers.
     class AdvancedPath:
@@ -127,6 +129,9 @@ def __init__(self, target: tuple) -> None:
         super().__init__(target)
         if not isinstance(target.arch, dict):
             raise TypeError("target.arch is not a dict")
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        mod = compile_module_from_src(Path(os.path.join(dirname, "arch_parser.c")).read_text(), "arch_utils")
+        self.parse_device_arch = mod.parse_device_arch
         self.properties = self.parse_target(target.arch)
         self.binary_ext = "spv"
 
@@ -142,30 +147,37 @@ def parse_target(self, tgt_prop) -> dict:
         dev_prop['max_num_sub_groups'] = tgt_prop.get('max_num_sub_groups', None)
         dev_prop['sub_group_sizes'] = tgt_prop.get('sub_group_sizes', None)
         dev_prop['has_fp64'] = tgt_prop.get('has_fp64', None)
-        if os.getenv("TRITON_INTEL_QUERY_DEVICE_EXTENSIONS", "0") == "1":
+        dev_prop['has_subgroup_matrix_multiply_accumulate'] = tgt_prop.get('has_subgroup_matrix_multiply_accumulate',
+                                                                           False)
+        dev_prop['has_subgroup_matrix_multiply_accumulate_tensor_float32'] = tgt_prop.get(
+            'has_subgroup_matrix_multiply_accumulate_tensor_float32', False)
+        dev_prop['has_subgroup_2d_block_io'] = tgt_prop.get('has_subgroup_2d_block_io', False)
+        dev_prop['has_bfloat16_conversions'] = tgt_prop.get('has_bfloat16_conversions', True)
+
+        device_arch = self.parse_device_arch(tgt_prop.get('architecture', 0))
+        if device_arch:
+            if device_arch in self.device_props:
+                dev_prop.update(self.device_props[device_arch])
+                return dev_prop
             try:
-                # FIXME: Add support for other devices.
-                ocloc_cmd = ['ocloc', 'query', 'CL_DEVICE_EXTENSIONS', '-device', 'pvc']
-                result = subprocess.run(ocloc_cmd, check=True, capture_output=True, text=True)
-                output = result.stdout
+                ocloc_cmd = ['ocloc', 'query', 'CL_DEVICE_EXTENSIONS', '-device', device_arch]
+                with tempfile.TemporaryDirectory() as temp_dir:
+                    output = subprocess.check_output(ocloc_cmd, text=True, cwd=temp_dir)
                 supported_extensions = set()
                 for extension in output.split(' '):
                     supported_extensions.add(extension)
-                dev_prop[
+                ocloc_dev_prop = {}
+                ocloc_dev_prop[
                     'has_subgroup_matrix_multiply_accumulate'] = 'cl_intel_subgroup_matrix_multiply_accumulate' in supported_extensions
-                dev_prop[
+                ocloc_dev_prop[
                     'has_subgroup_matrix_multiply_accumulate_tensor_float32'] = 'cl_intel_subgroup_matrix_multiply_accumulate_tensor_float32' in supported_extensions
-                dev_prop['has_subgroup_2d_block_io'] = 'cl_intel_subgroup_2d_block_io' in supported_extensions
-                dev_prop['has_bfloat16_conversions'] = 'cl_intel_bfloat16_conversions' in supported_extensions
-            except subprocess.CalledProcessError as e:
-                raise RuntimeError(f'`ocloc` failed with error code {e.returncode}')
-        else:
-            dev_prop['has_subgroup_matrix_multiply_accumulate'] = tgt_prop.get(
-                'has_subgroup_matrix_multiply_accumulate', False)
-            dev_prop['has_subgroup_matrix_multiply_accumulate_tensor_float32'] = tgt_prop.get(
-                'has_subgroup_matrix_multiply_accumulate_tensor_float32', False)
-            dev_prop['has_subgroup_2d_block_io'] = tgt_prop.get('has_subgroup_2d_block_io', False)
-            dev_prop['has_bfloat16_conversions'] = tgt_prop.get('has_bfloat16_conversions', True)
+                ocloc_dev_prop['has_subgroup_2d_block_io'] = 'cl_intel_subgroup_2d_block_io' in supported_extensions
+                ocloc_dev_prop['has_bfloat16_conversions'] = 'cl_intel_bfloat16_conversions' in supported_extensions
+                self.device_props[device_arch] = ocloc_dev_prop
+                dev_prop.update(ocloc_dev_prop)
+            except subprocess.CalledProcessError:
+                # Note: LTS driver does not support ocloc query CL_DEVICE_EXTENSIONS.
+                pass
         return dev_prop
 
     def parse_options(self, opts) -> Any: