iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.56k stars 568 forks source link

Remove topological sort from the `ScheduleExecution` pass. #7830

Open benvanik opened 2 years ago

benvanik commented 2 years ago

I don't recall why it's in there but it is seemingly load-bearing. I suspect the issue is that the execute op is inserted in the wrong place - the partition is a gather of ops within the block and picking the exact insertion location may not be possible (especially with clones) as-is.

Here's a minimal repro that fails with error: operand #2 does not dominate this use if sortBlockTopologically is commented out:

stream.executable private @ex {
  stream.executable.export public @dispatch_0
  stream.executable.export public @dispatch_1
  stream.executable.export public @dispatch_2
  builtin.module  {
    func @dispatch_0(%arg0: !stream.binding) {
      return
    }
    func @dispatch_1(%arg0: !stream.binding) {
      return
    }
    func @dispatch_2(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: !stream.binding) {
      return
    }
  }
}
func @predict(%arg0: i1) -> !stream.resource<transient> {
  %c1 = arith.constant 1 : index
  %c128 = arith.constant 128 : index
  %0 = stream.async.dispatch @ex::@dispatch_0[%c1, %c1, %c1]() : () -> !stream.resource<transient>{%c128}
  %1 = stream.async.dispatch @ex::@dispatch_1[%c1, %c1, %c1]() : () -> !stream.resource<transient>{%c128}
  %2 = select %arg0, %0, %1 : !stream.resource<transient>
  %3 = stream.async.dispatch @ex::@dispatch_2[%c1, %c1, %c1](%1, %2) : (!stream.resource<transient>{%c128}, !stream.resource<transient>{%c128}) -> !stream.resource<transient>{%c128}
  return %3 : !stream.resource<transient>
}
benvanik commented 2 years ago

Looks like an issue with hazard detection.