nod-ai / iree-amd-aie

IREE plugin repository for the AMD AIE accelerator
Apache License 2.0
46 stars 23 forks source link

Restore AIE/AIEX passes temporarily (2/n) #421

Closed makslevental closed 2 weeks ago

makslevental commented 2 weeks ago

This PR is part of a stack of PRs that refactor the dependency on MLIR-AIE. See https://github.com/nod-ai/iree-amd-aie/issues/430 for more information.

This PR adds all the boilerplate for the passes we'll ultimately need (and rewrite): aie-assign-lock-ids, aie-assign-buffer-descriptor-ids, aie-assign-buffer-addresses-basic, aie-pathfinder, aie-localize-locks, aie-objectstateful-transform, aiex-dma-to-npu. This PR also adds all the tests from upstream that'll serve as that unit-tests going forward.

Note, aie-assign-buffer-addresses-basic is based on the previous version of what mlir-aie currently has now and we take that one because it is dramatically simpler.

Also note, there's not much point in looking at the implementations of the passes closely until this PR https://github.com/nod-ai/iree-amd-aie/pull/427, at which point the implementation is the one that will be final-ish (i.e., up until that PR everything is mostly copy-pasted from mlir-aie, including bugs 😞).

makslevental commented 2 weeks ago

Eventually how many mlir dialects do you envisage here? amdaie & aie & aiex seems a bit unneccessary

Only amdaie. All analysis from aiex and aie will get folded into utility code (that manipulates non-AST representations).

Maybe create Passes.td generated to avoid boilerplate of aie/AIEAssignBufferAddressesBasic.cpp? That would also be a good place to add a sentence on what the pass does (as obvious as it should be now).

I didn't do this because the pass won't stay a pass - it is exactly the kind of simple utility code that has no need for all the boilerplate involved with transforming AST objects (it doesn't use any properties of the AST, it doesn't rewrite anything, it just updates an attribute that could be a simple field on a POD struct instead).

The pass AIEAssignBufferAddressesBasic is indeed nice and simple. But I guess we'll definitely want to make fine-grained decisions on which bank to allocate in eventually as in https://github.com/Xilinx/mlir-aie/blame/875648d47ede10c2405181325a325a91c7bad463/lib/Dialect/AIE/Transforms/AIEAssignBuffers.cpp thoughts?

Too complex but not complex enough (currently there are bugs). The right way to do this kind of thing is to write down a CPSAT problem and solve it. That was planned future work when I was in mlir-aie and now becomes planned future work here 🙂.

I don't know understand why you have added tests for passes which only exist in mlir-aie. I would bring the tests in if/when the pass being tested lives in iree-amd-aie.

Because they will serve as unit tests for the forthcoming PRs that vendor all the mlir-aie passes and refactor them. See https://github.com/nod-ai/iree-amd-aie/pull/425, https://github.com/nod-ai/iree-amd-aie/pull/426, https://github.com/nod-ai/iree-amd-aie/pull/427 (which aren't ready for review).