onnx / onnx-mlir

Representation and Reference Lowering of ONNX Models in MLIR Compiler Infrastructure
Apache License 2.0
767 stars 320 forks source link

Use top-down traversal in BundleMemoryPools #555

Open sstamenova opened 3 years ago

sstamenova commented 3 years ago

A recent (as of 3/20/21) change to mlir changed the traversal order from bottom up to top down. This breaks the algorithms in bundle memory pools. We get invalid outputs (see use of %8) such as:

  %0 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<10x20xf32>
  %1 = "krnl.getref"(%8, %c0_i64) : (memref<3200xi8>, i64) -> memref<10x20xf32>
  %2 = "krnl.getref"(%8, %c800_i64) : (memref<3200xi8>, i64) -> memref<10x10xf32>
  %3 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<400xi8>
  %4 = "krnl.getref"(%8, %c1200_i64) : (memref<3200xi8>, i64) -> memref<10x20xf32>
  %5 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<800xi8>
  %6 = "krnl.getref"(%8, %c2000_i64) : (memref<3200xi8>, i64) -> memref<10x20xf32>
  %7 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<800xi8>
  %8 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<3200xi8>
  %9 = "krnl.getref"(%8, %c2800_i64) : (memref<3200xi8>, i64) -> memref<10x10xf32>
  %10 = "memref.alloc"() {operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<400xi8>

The breaking change is: https://reviews.llvm.org/D99006 There is a temporary workaround checked in (but under discussion): https://reviews.llvm.org/D99059

At the end of the day, the change in traversal is meant to be permanent which means that onnx-mlir has to update the algorithms. I'm opening this as a heads up so that it's no surprise.

sstamenova commented 3 years ago

To make this even more interesting, the change was reverted: https://reviews.llvm.org/D99329 for the moment to give downstream components time to prepare.

See discussion here: https://llvm.discourse.group/t/speeding-up-canonicalize/3015/25

AlexandreEichenberger commented 3 years ago

@doru1004 any updates on this one?