Fix Flex Attn UT failures (#5138)

Before 92eb442, `getStride(ptr, 0)`
would be 1 for column major, i.e, pitch would be smaller than 64 bytes
(HW restriction), so `getPitch` would return `nullptr`, and 2d block
load would not be generated. This PR goes back to the original behavior
and early return for column major.

Flex Attn UT:
https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/17833144533/job/50703282622
(GOOD)
Flex Attn benchmark:
https://github.com/intel/intel-xpu-backend-for-triton/actions/runs/17836148975/job/50713713366
(GOOD)

Fixes #5129

---------

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

Author: whitneywhtsang

third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp

@@ -1867,9 +1867,12 @@ struct LoadOpToBlockIOConversion
     Type resultType = op.getType();
     auto tensorType = cast<RankedTensorType>(resultType);
 
-    // Step 1: Right now we only support 2D rank matrix of row major or column
-    // major.
+    // Step 1: Right now we only support 2D rank matrix of row major.
     const bool memoryRowMajor = isMemoryRowMajor(op);
+    // FIXME: Add support of column major.
+    if (!memoryRowMajor)
+      return failure();
+
     DpasEncodingAttr::OpIdx opIdx = getOpIdx(tensorType);
 
     Attribute encoding = tensorType.getEncoding();
@@ -1912,12 +1915,8 @@ struct LoadOpToBlockIOConversion
       // only support rank of 2 for now.
       return failure();
     }
-    const bool valueRowMajor =
-        (threadOrder[rank - 2] == 1 && threadOrder[rank - 1] == 0);
-    assert((valueRowMajor ||
-            (threadOrder[rank - 2] == 0 && threadOrder[rank - 1] == 1)) &&
-           "Only row_major or column_major is allowed");
-    const bool isTransposeRequired = valueRowMajor ^ memoryRowMajor;
+    unsigned contiguousDim = memoryRowMajor ? 1 : 0;
+    const bool isTransposeRequired = contiguousDim != colDim;
 
     // Step 2: Right now we only support DPAS related layout to simplify the
     // lowering.
@@ -2246,7 +2245,7 @@ struct LoadOpToBlockIOConversion
       return failure();
 
     // If the stride is 0, we want to load only the first row.
-    int stride = getStride(ptr, 0);
+    int stride = getStride(ptr, memoryRowMajor ? 0 : 1);
     unsigned baseHeightInt = (stride == 0 ? 1 : tileHeight);
     Value baseHeight = b.i32_val(baseHeightInt);
     Value baseWidth =