[Mlir-commits] [mlir] [MLIR] Extend linalg.pack and linalg.unpack to accept memref (PR #167675)

[Rolf Morel]

================
@@ -195,23 +194,12 @@ def Linalg_PackOp : Linalg_RelayoutOp<"pack", [
     //            expect tensor<2x8xf32> because CeilDiv(9, 8) = 2
     ```
   }];
-  let arguments = (ins AnyRankedTensor:$source,
-                       AnyRankedTensor:$dest,
-                       Optional<AnyType>:$padding_value,
-                       DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$outer_dims_perm,
-                       DenseI64ArrayAttr:$inner_dims_pos,
-                       Variadic<Index>:$inner_tiles,
-                       DenseI64ArrayAttr:$static_inner_tiles);
-  let results = (outs AnyRankedTensor:$result);
-  let assemblyFormat = [{
-    $source
-    (`padding_value` `(` $padding_value^ `:` type($padding_value) `)`)?
-    (`outer_dims_perm` `=` $outer_dims_perm^)?
-    `inner_dims_pos` `=` $inner_dims_pos
-    `inner_tiles` `=`
-    custom<DynamicIndexList>($inner_tiles, $static_inner_tiles)
-    `into` $dest attr-dict `:` type($source) `->` type($dest)
-  }];
+  let arguments = (ins TensorOrMemRef<[AnyType]>:$source,
+      TensorOrMemRef<[AnyType]>:$dest, Optional<AnyType>:$padding_value,
+      DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$outer_dims_perm,
+      DenseI64ArrayAttr:$inner_dims_pos, Variadic<Index>:$inner_tiles,
+      DenseI64ArrayAttr:$static_inner_tiles);
+  let results = (outs Optional<AnyRankedTensor>:$result);
----------------
rolfmorel wrote:

I think changing the op's formal results here is making the generated Python `PackOp` -- used [here](https://github.com/llvm/llvm-project/blob/f97f53e1d925fe702c23ad3a120ad3ab40b4eb51/mlir/python/mlir/dialects/linalg/__init__.py#L294-L355) to expose a nice `linalg.pack` Python wrapper -- take the result type as the first argument. You will just need a little bit of logic in this function (and presumably in the `linalg.unpack` Python wrapper as well) that determines if the op is using tensor semantics, in which case is should look something like `PackOp(dest.type, ...` and, in case of memref types/semantics, just `PackOp(None, ...)`. (N.B. syntax might be a bit different.)

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