nod-ai / iree-amd-aie

IREE plugin repository for the AMD AIE accelerator
Apache License 2.0
69 stars 30 forks source link

Use python bindings for tests #795

Closed makslevental closed 1 month ago

makslevental commented 2 months ago

This PR begins the migration to pytest for all of our e2e tests using IREE's python bindings for both the compiler and the runtime. I've tested it thoroughly (by running run_matmul.sh and cpu_comparison/run.py ~100~50 times).

Two major changes were required to make this happen:

  1. Move all CLI args to AMDAIEOptions so that they're registered at plugin load time. This was necessary for enabling Session.set_flags in Python. Going forward we should make sure to put all CLI args in the same place;
  2. The XRT HAL had to be drastically refactored because in its current state (prior to this PR), it doesn't support multiple dispatches in a single process - I had to track down all kinds of resource leaks (why this PR took so much longer than I thought);
    • There's a test 1536x2048x1536 that was prior disabled for Windows but is now disabled entirely. I don't know what the connection is (I wasn't able to track it down) but since every other test passes, including the 1000 run test, and even after repeating overall ~100~50 runs, I think it's safe to conclude there's something wrong with that shape rather than the HAL.

The migration isn't complete because the bindings and the VM itself don't support bf16 yet; I will send them a PR soon.

Note, in its current state the test scripts run ~100~50 times. I will remove this before landing.

TODO (for the next iteration): right now repeat_runs doesn't work because of how the XRT command buffer works (it's free'ed after every dispatch and with it the xrt::kernel and xrt::hw_context). Likely I'm doing something wrong (possibly using the wrong command buffer model) so I think it should be fixable.

newling commented 1 month ago

Off the top of my head things which we need to reach parity, hopefully we can reuse what's already done in cpu_comparison/run.py:

Also, can you/someone please at some point:

makslevental commented 1 month ago

keep track of migration of tests from old script to pytest, ideally they shouldn't be run in both places. Or maybe they've all been moved already?

None have been moved, some have been duplicated.