Merge remote-tracking branch 'origin/main' into etiotto.remove_masks

Author: etiotto

.github/workflows/build-test-reusable.yml

@@ -302,7 +302,7 @@ jobs:
         run: |
           echo "TRITON_TEST_CMD=${{ needs.build.outputs.test-triton-command }}" | tee -a $GITHUB_ENV
 
-      - name: Build PTI && Run Proton tests
+      - name: Run Proton tests
         if: matrix.suite == 'rest'
         run: |
           export LD_LIBRARY_PATH=${{ env.PTI_LIBS_DIR }}:$LD_LIBRARY_PATH

.github/workflows/triton-benchmarks.yml

@@ -117,17 +117,10 @@ jobs:
           cd benchmarks
           pip install .
 
-      - name: Build PTI from source
+      - name: Build PTI
         run: |
-          PTI_COMMIT_ID="$(<.github/pins/pti.txt)"
-          git clone https://github.com/intel/pti-gpu.git
-          cd pti-gpu
-          git checkout $PTI_COMMIT_ID
-          cd sdk
-          cmake --preset linux-icpx-release
-          BUILD_TESTING=1 PTI_BUILD_SAMPLES=1 cmake --build --preset linux-icpx-release
-
-          PTI_LIBS_DIR="$(pwd)/build-linux-icpx-release/lib/"
+          ./scripts/install-pti.sh --build-level-zero
+          PTI_LIBS_DIR=$(python ./scripts/pti_lib.py)
           ls $PTI_LIBS_DIR
           echo "PTI_LIBS_DIR=$PTI_LIBS_DIR" >> $GITHUB_ENV
 

python/test/unit/tools/test_disasm.py

@@ -34,4 +34,4 @@ def kernel(X, i: tl.constexpr):
     assert x[0] == 12
     dis = h.asm["spvdis"]
     # check that the spvdis has a store instruction.
-    assert "PredicatedStore" in dis
+    assert "OpStore" in dis

python/triton/compiler/compiler.py

@@ -468,6 +468,12 @@ def raise_(err):
         self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = driver.active.utils.load_binary(
             self.name, self.kernel, self.metadata.shared, self.metadata.build_flags,
             not self.metadata.generate_native_code, device)
+        # PyTorch could use the updated build flags in load binary.
+        if hasattr(driver.active.utils, "get_last_selected_build_flags"):
+            new_build_flags = driver.active.utils.get_last_selected_build_flags()
+            if new_build_flags != self.metadata.build_flags:
+                self.metadata = self.metadata._replace(build_flags=new_build_flags)
+
         if hasattr(self.metadata, "threads_per_warp"):
             warp_size = self.metadata.threads_per_warp
         else:

third_party/intel/backend/driver.c

@@ -192,6 +192,12 @@ sycl::context get_default_context(const sycl::device &sycl_device) {
 #endif
 }
 
+static BuildFlags last_build_flag("");
+
+extern "C" EXPORT_FUNC PyObject *get_last_selected_build_flags() {
+  return Py_BuildValue("s", last_build_flag().data());
+}
+
 extern "C" EXPORT_FUNC PyObject *load_binary(PyObject *args) {
   const char *name, *build_flags_ptr;
   int shared;
@@ -309,7 +315,7 @@ extern "C" EXPORT_FUNC PyObject *load_binary(PyObject *args) {
         PyCapsule_New(reinterpret_cast<void *>(fun), "kernel", freeKernel);
     auto kernel_bundle_py = PyCapsule_New(reinterpret_cast<void *>(mod),
                                           "kernel_bundle", freeKernelBundle);
-
+    last_build_flag = build_flags;
     return Py_BuildValue("(OOiii)", kernel_bundle_py, kernel_py, n_regs,
                          n_spills, n_max_threads);
 

third_party/intel/backend/driver.py

@@ -199,9 +199,11 @@ def __init__(self, cache_path: str):
         self.shared_library.get_device_properties.argtypes = (ctypes.c_int, )
         self.shared_library.has_opencl_extension.restype = ctypes.py_object
         self.shared_library.has_opencl_extension.argtypes = (ctypes.c_int, ctypes.c_char_p)
+        self.shared_library.get_last_selected_build_flags.restype = ctypes.py_object
 
     def __getattribute__(self, name):
-        if name in ("get_device_properties", "init_devices", "wait_on_sycl_queue", "has_opencl_extension"):
+        if name in ("get_device_properties", "init_devices", "wait_on_sycl_queue", "has_opencl_extension",
+                    "get_last_selected_build_flags"):
             shared_library = super().__getattribute__("shared_library")
             return getattr(shared_library, name)
 
@@ -318,6 +320,7 @@ def __init__(self):
         self.device_count = mod.init_devices(self.get_sycl_queue())
         self.wait_on_sycl_queue = mod.wait_on_sycl_queue
         self.has_opencl_extension = mod.has_opencl_extension
+        self.get_last_selected_build_flags = mod.get_last_selected_build_flags
 
     def get_current_device(self):
         import torch

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -3401,10 +3401,7 @@ struct StoreOpConversion
 
       if (maskVal) {
         // Create a predicated store operation.
-        std::optional<bool> enablePredicated =
-            mlir::triton::tools::isEnvValueBool(
-                mlir::triton::tools::getStrEnv("TRITON_INTEL_PREDICATED"));
-        if (!enablePredicated.has_value() || enablePredicated.value())
+        if (triton::tools::getBoolEnv("TRITON_INTEL_PREDICATED"))
           rewriter.create<TritonGEN::PredicatedStoreOp>(
               loc, addrElem, vecWord, b.i64_val(alignment), maskVal);
         else