[Mlir-commits] [mlir] [mlir][linalg] Add PackOp canonicalization pattern (PR #215785)

[Hendrik Klug]

================
@@ -6008,6 +6008,13 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
     }
   }
 
+  // Fold pack(empty) to the destination tensor if no padding value is provided.
+  if (packOp.getSource().getDefiningOp<tensor::EmptyOp>() &&
----------------
HendrikKlug-synthara wrote:

hmm I don't think the resulting IR is broken, the original IR fills an empty into %1 and returns that.
It also packs the empty into some provided %arg1 and returns the packed data.

After canonicalisation it still returns the filled empty %1 but instead of packing an empty into %arg1, it directly returns %arg1.
The IR changes by returning %arg1 data instead of unspecified data, can that break anything?

For the case that you describe the folding already exists in main, where the empty gets filled with %arg0, packed (argument of pack is the output of the fill op, not a tensor.empty) and returned:
instead of filling a packed empty, after canonicalisation %arg1 gets filled directly and returned, skipping the pack.
Bufferization creates an extra allocation for the "filled %arg1" so that should be safe.

```mlir

module {
  func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> tensor<4x8x8x32xf32> {
    %0 = tensor.empty() : tensor<64x128xf32>
    %1 = linalg.fill ins(%arg0 : f32) outs(%0 : tensor<64x128xf32>) -> tensor<64x128xf32>
    %pack = linalg.pack %1 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %arg1 : tensor<64x128xf32> -> tensor<4x8x8x32xf32>
    return %pack : tensor<4x8x8x32xf32>
  }
}

// -----// IR Dump After CanonicalizerPass: canonicalize{cse-between-iterations=false    max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true} //----- //
module {
  func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> tensor<4x8x8x32xf32> {
    %0 = linalg.fill ins(%arg0 : f32) outs(%arg1 : tensor<4x8x8x32xf32>) -> tensor<4x8x8x32xf32>
    return %0 : tensor<4x8x8x32xf32>
  }
}


// -----// IR Dump After OneShotBufferizePass: one-shot-bufferize{allow-return-allocs-from-loops=false allow-unknown-ops=false analysis-fuzzer-seed=0 analysis-heuristic=bottom-up buffer-alignment=64 bufferize-function-boundaries=false check-parallel-regions=true copy-before-write=false  dump-alias-sets=false function-boundary-type-conversion=infer-layout-map must-infer-memory-space=false  print-conflicts=false test-analysis-only=false unknown-type-conversion=fully-dynamic-layout-map use-encoding-for-memory-space=false} //----- //
module {
  func.func @write_then_pack(%arg0: f32, %arg1: tensor<4x8x8x32xf32>) -> tensor<4x8x8x32xf32> {
    %alloc = memref.alloc() {alignment = 64 : i64} : memref<4x8x8x32xf32>
    linalg.fill ins(%arg0 : f32) outs(%alloc : memref<4x8x8x32xf32>)
    %0 = bufferization.to_tensor %alloc : memref<4x8x8x32xf32> to tensor<4x8x8x32xf32>
    return %0 : tensor<4x8x8x32xf32>
  }
}


```





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