nod-ai / iree-amd-aie

IREE plugin repository for the AMD AIE accelerator
Apache License 2.0
64 stars 29 forks source link

[AMD-AIE] Fix single iteration promotion for scf::ForallOp #381

Closed Abhishek-Varma closed 4 months ago

Abhishek-Varma commented 4 months ago

-- This commit adds a fix for single iteration promotion for scf::ForallOp.

Signed-off-by: Abhishek Varma abhvarma@amd.com

Repro for the issue:

  1. You may take either the lit test as the INPUT IR or controlcode_input.mlir which we get via the C++ pipeline for objectfifo backend.
  2. Run iree-opt --pass-pipeline="builtin.module(func.func(iree-amdaie-controlcode-loop-unroll,canonicalize))" input.mlir.

The promotion of the single iteration Forall spuriously shifts the body of the loop without the SSA values replaced.

Surprisingly the same input IR (lit test at least) is targeted for using transform dialect - IT WORKS.

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    transform.loop.promote_if_one_iteration %0 : !transform.any_op
    transform.yield
  }
}

You may use the above to confirm.

I tried triaging and found https://github.com/llvm/llvm-project/blob/8eb0945373173213e7454a475f6e227da12d6d3a/mlir/lib/Dialect/SCF/IR/SCF.cpp#L650 to be the culprit. It creates constants for LB correctly when using Transform dialect script, but fails to do that when invoking the same function (which even the Transform dialect does).

Perhaps Transform interpreter is already taking care of some insertion logic which undoes the spurious effect?

Seems like an upstream issue - I'll raise a PR upstream too but the reproducibility can be proven by a test pass as the same function passes via Transform dialect.