Reland de1f346aa6737fa2e3e6a8a64dae118fcfab9995 that is temporarily reverted in https://github.com/intel/intel-xpu-backend-for-triton/pull/2793 by 7b5daa47fd38226587c1e6be35d6e6e371074735.