[Mlir-commits] [mlir] [mlir][tensor] Add e2e test for tensor.unpack with dynamic tile sizes (PR #121557)

[Han-Chung Wang]

================
@@ -0,0 +1,110 @@
+// DEFINE: %{compile} =  mlir-opt %s \
+// DEFINE:  -transform-interpreter -test-transform-dialect-erase-schedule |\
+// DEFINE: mlir-opt \
+// DEFINE:  -test-lower-to-llvm -o %t
+// DEFINE: %{entry_point} = main
+// DEFINE: %{run} = mlir-cpu-runner %t -e %{entry_point} -entry-point-result=void \
+// DEFINE:    -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
+
+// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s
+
+/// End-to-end test for tensor.unpack where one of the inner tile sizes is
+/// dynamic. See pack-dynamic-inner-tile.mlir for a similar test for tensor.pack.
+
+func.func @main() {
+  // Allocate and initialise the inputs
+  %A_alloc = tensor.empty() : tensor<7x3xi32>
+
+  %A = arith.constant dense<[
+  [[[1],
+   [2],
+   [3],
+   [4],
+   [5],
+   [6],
+   [7],
+   [123]],
+  [[8],
+   [9],
+   [10],
+   [11],
+   [12],
+   [13],
+   [14],
+   [123]],
+  [[15],
+   [16],
+   [17],
+   [18],
+   [19],
+   [20],
+   [21],
+   [123]]]
+  ]> : tensor<1x3x8x1xi32>
+
+  %A_cast = tensor.cast %A : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
+  func.call @unpack(%A_cast) : (tensor<?x3x?x1xi32>) -> ()
+
+  return
+}
+
+func.func private @unpack(%A: tensor<?x3x?x1xi32>) {
+  %c1 = arith.constant 1 : index
+  %pad_val = arith.constant 123 : i32
+
+  // Dynamic tile size
+  %tile_size = arith.constant 8 : index
----------------
hanhanW wrote:

I think you are right about the scalability part.

> Now, do we need to worry about this though? The test specifies it's own lowering pipeline (through TD) and canonicalization is used fairly late. So perhaps it will be fine?

I think we do not need to worry about it for now. It is mostly just a note. Some drivers could kick in folders (e.g., OpBuilder::createOrFold) and it becomes an issue when people add folding methods to the op: https://mlir.llvm.org/docs/Canonicalization/#canonicalizing-with-the-fold-method

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