[Mlir-commits] [mlir] [mlir][linalg] Support scalable vectorization of linalg.index operations (PR #96778)

[Andrzej Warzyński]

================
@@ -113,6 +113,56 @@ module attributes {transform.with_named_sequence} {
   }
 }
 
+// -----
+
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func.func @vectorize_scalable_nd_tensor_extract_transfer_read_basic(%arg0: tensor<?x?x?xf32>, %arg2: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %1 = linalg.generic {
+    indexing_maps = [#map1],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } outs(%arg2 : tensor<?x?x?xf32>) {
+  ^bb0(%arg4: f32):
+    %2 = linalg.index 0 : index
+    %3 = linalg.index 1 : index
+    %4 = linalg.index 2 : index
+    %5 = tensor.extract %arg0[%2, %3, %4] : tensor<?x?x?xf32>
+    linalg.yield %5 : f32
+  } -> tensor<?x?x?xf32>
----------------
banach-space wrote:

Only the trailing dim is required to be dynamic to test this contribution? Let's keep things simple.

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