Add knobs TRITON_INTEL_ENABLE_BLOCK_IO_STORE_ON_REGULAR_PTR for block store. (#4610)

The 2D block IO store causes bad performance on flex attn. 
Disable the block store by default and add a knob to enable it.

---------

Signed-off-by: Lu,Chengjun <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: chengjunlu

include/triton/Tools/Sys/GetEnv.hpp

@@ -51,6 +51,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_INTEL_ENABLE_INSTR_SCHED",
     "TRITON_INTEL_FAST_MATH",
     "TRITON_INTEL_REDUCE_TRANSPOSE",
+    "TRITON_INTEL_ENABLE_BLOCK_IO_STORE_ON_REGULAR_PTR",
     // clang-format on
 };
 

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -14,6 +14,7 @@
 #include "intel/include/Dialect/TritonIntelGPU/IR/Attributes.h"
 #include "intel/include/Dialect/TritonIntelGPU/Transforms/Utility.h"
 #include "triton/Tools/LinearLayout.h"
+#include <triton/Tools/Sys/GetEnv.hpp>
 
 using namespace mlir;
 using namespace mlir::triton;
@@ -2686,6 +2687,11 @@ struct StoreOpToBlockIOConversion
       offsetBaseX = b.trunc(i32_ty, offsetX);
       offsetBaseY = b.trunc(i32_ty, offsetY);
     } else {
+
+      static const bool enableBlockStore = triton::tools::getBoolEnv(
+          "TRITON_INTEL_ENABLE_BLOCK_IO_STORE_ON_REGULAR_PTR");
+      if (!enableBlockStore)
+        return failure();
       // Get the LLVM values for pointers
       ptrElems = unpackLLElements(loc, llPtr, rewriter);
       assert(ptrElems.size() == numElems &&