[Mlir-commits] [mlir] [mlir][XeGPU][Transform] Add gather/scatter coalescing analysis. (PR #201684)

[Charitha Saumya]

================
@@ -0,0 +1,257 @@
+// RUN: mlir-opt -split-input-file \
+// RUN:   -test-xegpu-coalesce-gather-scatter="analyze-only=true" %s | FileCheck %s
+
+// Contiguity analysis: stamps the `contiguity` attribute on gather/scatter ops
+// whose offsets are contiguous (runs of >= 2) along the innermost dimension.
+// The stamped value is the inner-dim contiguity, rounded down to a divisor of
+// the inner extent. The analysis is target-independent and mask-independent;
+// the lane_data split and any access-pattern gating live in the consumer.
+
+// -----
+// 1-D vector.step -> stride-1, fully contiguous over the 32-element inner dim.
+// CHECK-LABEL: func.func @load_step_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_step_offsets(%ptr: i64) -> vector<32xf32> {
+  %offsets = vector.step : vector<32xindex>
+  %mask = arith.constant dense<true> : vector<32xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+  return %v : vector<32xf32>
+}
+
+// -----
+// Dense stride-1 constant offsets -> contiguity 32.
+// CHECK-LABEL: func.func @load_dense_ap_offsets(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_dense_ap_offsets(%ptr: i64) -> vector<32xi32> {
+  %offsets = arith.constant dense<[
+    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
+    16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
+  ]> : vector<32xindex>
+  %mask = arith.constant dense<true> : vector<32xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<32xindex>, vector<32xi1> -> vector<32xi32>
+  return %v : vector<32xi32>
+}
+
+// -----
+// Stride-4 offsets: not contiguous, no attribute stamped.
+// CHECK-LABEL: func.func @load_stride4_no_attr(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+// CHECK-SAME: : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+func.func @load_stride4_no_attr(%ptr: i64) -> vector<32xf32> {
+  %c4 = arith.constant 4 : index
+  %step = vector.step : vector<32xindex>
+  %splat = vector.broadcast %c4 : index to vector<32xindex>
+  %offsets = arith.muli %step, %splat : vector<32xindex>
+  %mask = arith.constant dense<true> : vector<32xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+  return %v : vector<32xf32>
+}
+
+// -----
+// All-equal offsets: inner dim is constant, not contiguous -> no attribute.
+// CHECK-LABEL: func.func @load_broadcast_offsets_no_attr(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_broadcast_offsets_no_attr(%ptr: i64) -> vector<32xf32> {
+  %offsets = arith.constant dense<0> : vector<32xindex>
+  %mask = arith.constant dense<true> : vector<32xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+  return %v : vector<32xf32>
+}
+
+// -----
+// The analysis is mask-independent: a partial mask still gets the attribute
+// (the mask check is a consumer concern).
+// CHECK-LABEL: func.func @load_partial_mask(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_partial_mask(%ptr: i64, %mask: vector<32xi1>) -> vector<32xf32> {
+  %offsets = vector.step : vector<32xindex>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+  return %v : vector<32xf32>
+}
+
+// -----
+// Store with vector.step offsets -> contiguity 32 on the store.
+// CHECK-LABEL: func.func @store_step_offsets(
+// CHECK: xegpu.store
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @store_step_offsets(%ptr: i64, %v: vector<32xf32>) {
+  %offsets = vector.step : vector<32xindex>
+  %mask = arith.constant dense<true> : vector<32xi1>
+  xegpu.store %v, %ptr[%offsets], %mask
+      : vector<32xf32>, i64, vector<32xindex>, vector<32xi1>
+  return
+}
+
+// -----
+// 2-D leading-unit dim: contiguity measured on the inner dim -> 32.
+// CHECK-LABEL: func.func @load_2d_leading_unit(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_2d_leading_unit(%ptr: i64) -> vector<1x32xf32> {
+  %step = vector.step : vector<32xindex>
+  %offsets = vector.shape_cast %step : vector<32xindex> to vector<1x32xindex>
+  %mask = arith.constant dense<true> : vector<1x32xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<1x32xindex>, vector<1x32xi1> -> vector<1x32xf32>
+  return %v : vector<1x32xf32>
+}
+
+// -----
+// True 2-D dense AP: each row stride-1 over 16 -> contiguity 16.
+// CHECK-LABEL: func.func @load_2d_dense_ap(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 16 : i64}>
+func.func @load_2d_dense_ap(%ptr: i64) -> vector<2x16xf32> {
+  %offsets = arith.constant dense<[
+    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
+    [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+  ]> : vector<2x16xindex>
+  %mask = arith.constant dense<true> : vector<2x16xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+  return %v : vector<2x16xf32>
+}
+
+// -----
+// 2-D dense values whose inner row is not a stride-1 AP: no attribute.
+// CHECK-LABEL: func.func @load_2d_non_ap(
+// CHECK: xegpu.load
+// CHECK-NOT: contiguity
+func.func @load_2d_non_ap(%ptr: i64) -> vector<2x16xf32> {
+  %offsets = arith.constant dense<[
+    [0, 1, 2, 3, 4, 5, 6, 7, 100, 9, 10, 11, 12, 13, 14, 15],
+    [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+  ]> : vector<2x16xindex>
+  %mask = arith.constant dense<true> : vector<2x16xi1>
+  %v = xegpu.load %ptr[%offsets], %mask
+      : i64, vector<2x16xindex>, vector<2x16xi1> -> vector<2x16xf32>
+  return %v : vector<2x16xf32>
+}
+
+// -----
+// `divui` by a constant equal to the inner stride recovers stride-1
+// contiguity: [0,2,..,62] / 2 = [0,1,..,31] -> contiguity 32.
+// CHECK-LABEL: func.func @load_divui_recovers(
+// CHECK: xegpu.load
+// CHECK-SAME: <{contiguity = 32 : i64}>
+func.func @load_divui_recovers(%ptr: i64) -> vector<32xf32> {
+  %even = arith.constant dense<[
+     0,  2,  4,  6,  8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30,
+    32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62]>
+      : vector<32xindex>
+  %c2 = arith.constant dense<2> : vector<32xindex>
+  %offsets = arith.divui %even, %c2 : vector<32xindex>
----------------
charithaintc wrote:

nit: interesting test. if we add `SparseConstantPropagation` to the base analysis of the solver I think this div would be folded in the analysis. so I am not quite sure if we should handle arith op cases where all operands are constants. thoughts?

https://github.com/llvm/llvm-project/pull/201684