[Mlir-commits] [mlir] [mlir][linalg] Fix vectorizer generating invalid vector.gather for 0-D tensor.extract (PR #187085)

[Andrzej Warzyński]

================
@@ -477,3 +477,47 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// Rank-0 linalg.generic with tensor.extract using a data-dependent index.
+// The tensor.extract should be classified as ScalarBroadcast (not Gather),
+// producing a vector.transfer_read of a 0-D vector.
+
+func.func @rank0_tensor_extract_data_dependent_index(
+    %src: tensor<2xi64>,
+    %idx_tensor: tensor<i64>) -> tensor<i64> {
+
+  %init = tensor.empty() : tensor<i64>
+  %res = linalg.generic {
+    indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>],
+    iterator_types = []}
+    ins(%idx_tensor : tensor<i64>) outs(%init : tensor<i64>) {
+  ^bb0(%in: i64, %out: i64):
+    %idx = arith.index_cast %in : i64 to index
+    %val = tensor.extract %src[%idx] : tensor<2xi64>
+    linalg.yield %val : i64
+  } -> tensor<i64>
+
+  return %res : tensor<i64>
+}
+
+// CHECK-LABEL: func.func @rank0_tensor_extract_data_dependent_index(
+// CHECK-SAME:      %[[SRC:.*]]: tensor<2xi64>,
+// CHECK-SAME:      %[[IDX_TENSOR:.*]]: tensor<i64>) -> tensor<i64> {
+// CHECK-DAG:       %[[INIT:.*]] = tensor.empty() : tensor<i64>
+// CHECK-DAG:       %[[PAD:.*]] = ub.poison : i64
+// CHECK:           %[[READ_IDX:.*]] = vector.transfer_read %[[IDX_TENSOR]][], %[[PAD]] : tensor<i64>, vector<i64>
+// CHECK:           %[[SCALAR_IDX:.*]] = vector.extract %[[READ_IDX]][] : i64 from vector<i64>
+// CHECK:           %[[INDEX:.*]] = arith.index_cast %[[SCALAR_IDX]] : i64 to index
+// CHECK:           %[[READ_VAL:.*]] = vector.transfer_read %[[SRC]][%[[INDEX]]], %{{.*}} : tensor<2xi64>, vector<i64>
+// CHECK:           %[[WRITE:.*]] = vector.transfer_write %[[READ_VAL]], %[[INIT]][] : vector<i64>, tensor<i64>
+// CHECK:           return %[[WRITE]] : tensor<i64>
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.generic"]} in %module : (!transform.any_op) -> !transform.any_op
+    transform.structured.vectorize %0 {vectorize_nd_extract} : !transform.any_op
----------------
banach-space wrote:

If you don't care about vector sizes (makes sense), then could you move this to https://github.com/llvm/llvm-project/blob/ed76cbccbb40da23f443ee147d697917c5101e9a/mlir/test/Dialect/Linalg/vectorization/extract-with-patterns.mlir instead? Thanks!

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