[Mlir-commits] [mlir] [mlir][linalg] Collapse ND tensor.extract sources before gathering (PR #237308)

[Federico Bruzzone]

================
@@ -526,3 +534,122 @@ module attributes {transform.with_named_sequence} {
      transform.yield
    }
 }
+
+// -----
+
+// Only the inactive fourth lane is out of bounds: the linear offsets are
+// [1, 3, 5, 7], while the collapsed source contains six elements.
+func.func @masked_gather_inactive_lane_out_of_bounds(
+    %src: tensor<3x2xf32>, %out: tensor<3xf32>) -> tensor<3xf32> {
+  %c1 = arith.constant 1 : index
+  %result = linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  } outs(%out : tensor<3xf32>) {
+  ^bb0(%unused: f32):
+    %i = linalg.index 0 : index
+    %value = tensor.extract %src[%i, %c1] : tensor<3x2xf32>
+    linalg.yield %value : f32
+  } -> tensor<3xf32>
+  return %result : tensor<3xf32>
+}
+
+// CHECK-LABEL: func.func @masked_gather_inactive_lane_out_of_bounds(
+// CHECK-SAME: %[[SRC:.*]]: tensor<3x2xf32>,
+// CHECK-DAG: %[[THREE:.*]] = arith.constant 3 : index
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[THREE]] : vector<4xi1>
+// CHECK: %[[STEP:.*]] = vector.step : vector<4xindex>
+// CHECK-DAG: %[[ALL_TRUE:.*]] = arith.constant dense<true> : vector<4xi1>
+// CHECK-DAG: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
+// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : index
+// CHECK: %[[INNER_DIM:.*]] = arith.constant 1 : index
+// CHECK: %[[DIM:.*]] = tensor.dim %[[SRC]], %[[INNER_DIM]] : tensor<3x2xf32>
+// CHECK: %[[STRIDE:.*]] = vector.broadcast %[[DIM]] : index to vector<4xindex>
+// CHECK: %[[ROW:.*]] = arith.muli %[[STEP]], %[[STRIDE]] : vector<4xindex>
+// CHECK: %[[COL:.*]] = arith.constant dense<1> : vector<4xindex>
+// CHECK: %[[OFFSETS:.*]] = arith.addi %[[COL]], %[[ROW]] : vector<4xindex>
+// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape %[[SRC]] {{\[\[}}0, 1]] : tensor<3x2xf32> into tensor<6xf32>
+// CHECK: %[[GATHER:.*]] = vector.mask %[[MASK]] { vector.gather %[[COLLAPSED]][%[[ZERO]]] [%[[OFFSETS]]], %[[ALL_TRUE]], %[[PASS_THRU]] : tensor<6xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[GATHER]],
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%root: !transform.any_op {transform.readonly}) {
+    %target = transform.structured.match ops{["linalg.generic"]} in %root
+      : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %target vector_sizes [4] vectorize_nd_extract
+      : !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// Each source dimension fits in int64, but the flattened extent does not.
+func.func @gather_source_element_count_overflow(
+    %src: tensor<3037000500x3037000500xf32>, %out: tensor<2xf32>) -> tensor<2xf32> {
+  %c0 = arith.constant 0 : index
+  %result = linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  } outs(%out : tensor<2xf32>) {
+  ^bb0(%unused: f32):
+    %i = linalg.index 0 : index
+    %value = tensor.extract %src[%i, %c0] : tensor<3037000500x3037000500xf32>
+    linalg.yield %value : f32
+  } -> tensor<2xf32>
+  return %result : tensor<2xf32>
+}
+
+// CHECK-LABEL: func.func @gather_source_element_count_overflow(
+// CHECK-NOT: vector.
----------------
FedericoBruzzone wrote:

`CHECK-NOT: vector.` before and after, plus `CHECK: return`, is too broad from my POV

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