Allow transforms to be added as Python files

[rengolin]

This allows for a schedule generator in Python to be parametrized by a dictionary of options (similar to passes) that generates an ir.Module with a schedule inside. The main difference from importing from MLIR files is that you can parametrize the generation (similar to C++ templates).

Like LLVM and MLIR plugins, this allows code external to the library to execute, but that's true for any python file that uses the library or shared objects that are loaded, so no additional security issues here. As long as the code returns a valid MLIR module with a schedule inside, the schedule can be applied to the payload.

With this option we can start using the various schedules that are being added to lighthouse directly, but keeping the door open for user-defined schedules as well. We don't want to force users of this driver to have to register their own local transforms via some python mechanism, especially if they're not creating a python wrapper in the first place (ex. using lh-opt or a compiler).

[adam-smnk]

I'm thinking how this will interact with lighthouse's collection of schedules
Current opt scheme would require fixed API and creating separate file for each schedule.
Doable but very rigid approach which further hard couples opt and schedule modules instead of loose composition.

While I'd prefer something like registry mechanism, it might not be worth spinning up whole new infrastructure just yet.
As a middle ground, how about requiring to pass both path to a file to load and a name of schedule to use?
For example: lh-opt --stage=path/to/file.py[name].
There could even be further syntax to pass (trivial) parameters: [tile](tile_size=32) etc.

That would allow to have files (submodules) containing multiple schedules with separate APIs for those that want to use it as a pure library. For opt (CLI and Stages framework), it'd require only limited lazy symbol search which should minimize complexity and overhead.

[rengolin]

As a middle ground, how about requiring to pass both path to a file to load and a name of schedule to use?
For example: lh-opt --stage=path/to/file.py[name].
There could even be further syntax to pass (trivial) parameters: [tile](tile_size=32) etc.

Was discussing this with @rolfmorel just now. We reached a similar conclusion. Though, I'd still have a "default" naming scheme just for the sake of brevity. This PR is the "default" mechanism. I'll add the extension to a follow up PR.

[adam-smnk]

Just for my understanding, there's no way to pass parameters to the schedule yet?

[rengolin]

Correct. This will be done in follow-up PRs.

I like your syntax above, and unifying for passes, schedules and python files:

  # This is already available
  --stage=func.func(passname{opt1=val1 opt2=val2})

  # This may require changes to the upstream API
  --stage=schedule.mlir[sched_name]{opt1=val1 opt2=val2}

  # This may need differentiation between python argument and schedule arguments
  # The latter will require the same changes as the one above
  --stage=schedule.py[func_name,sched_name]{opt1=val1}{opt2=val2}

With func_name and sched_name having default values, which are the ones in this PR.

@rolfmorel also suggested we could have a python file that just implements the Stage interface, but it's not clear what the syntax would be. We can look at it later.