[Mlir-commits] [mlir] [mlir][Linalg] Allow more control in drop unit dims (PR #170104)

[Andrzej Warzyński]

================
@@ -24,3 +25,44 @@ func.func @drop_outermost_unit_dims(%arg0: tensor<1x1x42xf32>) -> tensor<1x1x42x
 //  CHECK-SAME:       outs(%[[OUTS_RESHAPE]] :
 //       CHECK:   %[[EXPAND_SHAPE:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0, 1], [2]{{\]}}
 //       CHECK:   return %[[EXPAND_SHAPE]]
+
+// -----
+
+#encoding = #test.tensor_encoding<"encoding">
+
+// Test that tensor encodings are preserved when collapsing unit dimensions
+func.func @drop_outermost_unit_dims_with_encoding(%arg0: tensor<1x1x42xf32, #encoding>) -> tensor<1x1x42xf32, #encoding> {
+  %0 = tensor.empty() : tensor<1x1x42xf32, #encoding>
+  %1 = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+                     affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+    iterator_types = ["parallel", "parallel", "parallel"]}
+    ins(%arg0 : tensor<1x1x42xf32, #encoding>)
+    outs(%0 : tensor<1x1x42xf32, #encoding>) {
+      ^bb0(%b0: f32, %b1 : f32):
+        %2 = arith.addf %b0, %b1 : f32
+        linalg.yield %2 : f32
+    } -> tensor<1x1x42xf32, #encoding>
+  return %1 : tensor<1x1x42xf32, #encoding>
+}
----------------
banach-space wrote:

Isn't this input identical to the one above? Why repeat it? As in, is a dedicated example with `#encoding` required?

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