Generic workload object and execution model

[tkarna]

We need a generic mechanism for defining new workloads and executing them through an unified execution engine interface.

With such an interface, it should be easy to get workload performance metrics, define benchmark suites, hook up workloads to CI and/or an autotuner, etc.

A workload is defined by the specific payload kind (e.g., a matmul, or MLP) with a fixed problem size. For the time being, we can assume that all input shapes are static, implying that the payload IR is fixed. A workload may involve an arbitrary lowering schedule (potentially with parameters, e.g., for tile sizes) and target arbitrary hardware. As such, they may require custom tools, such as custom methods to allocate/deallocate the payload input arrays before/after execution.

The Workload object could define methods:

• Constructor: Sets problem size and other payload options (e.g., elementwise post-processing ops). These are fixed.
• payload_module(): Returns a new module with the payload function and necessary MLIR utility functions (if any). This module will be lowered in-place.
• schedule_module(parameters): Returns a new transform schedule module. The provided parameters sets all optional schedule parameters (if any).
• get_input_arrays(execution_engine): Returns handles to payload function input arrays. The execution engine is provided in case MLIR helper functions need to be called.
• allocate(execution_engine): A context manager that does necessary allocation/deallocation of resources for execution (if any).
• verify(execution_engine): Runs correctness test, e.g., against a reference solution. The execution engine is provided in case MLIR helper functions need to be called.
• requirements(): Returns a list of ExecutionEngine requirements, e.g. MLIR libs or GPU runtime lib.
• get_complexity(): Returns the computational complexity of the workload, e.g. number of floating points operations and number of bytes read/written from/to memory.

The execution engine interface could provide high-level methods such as:

• lower_payload(workload, schedule_parameters): Returns the payload IR lowered using the transform schedule and its given parameters.
• execute(workload, schedule_parameters, check_correctness=True): Lowers the payload IR, allocates resources, executes it, (optionally) checks correctness, and deallocates resources.
• benchmark(workload, schedule_parameters, check_correctness=True): Executes the workload in a timer loop and returns measured timings. The get_complexity method can be used to compute throughput etc metrics.

One useful property is that the same Workload object can be re-used between execution calls, and therefore we can optionally cache some data in the object. For example, in some cases it's possible to cache the payload input arrays and the reference solution, which speeds up autotuning workflows, for example, where the same workload may be executed thousands of times.

Addresses point 2. of #15.

[tkarna]

For simple workloads, e.g., NumPy-expressible ops with standard data types like float32 running on the host, it is possible to a) allocate input arrays, b) initialize them, c) copy arrays to/from device (no-op), d) compute reference solution, and e) deallocate the arrays in Python e.g. with Numpy. In more complex cases, e.g., using more esoteric data types (Float8E5M2) in a GPU kernel it might be necessary to emit MLIR helper functions for all of the 5 steps (a-e).

[tkarna]

The above execute and benchmark methods are just examples of common use cases. Using the Workload methods and execution engine utility functions (not detailed here), developers can write their own custom kernel execution logic if they want to.

[tkarna]

Targeting different hardware implies using a different transform schedule and potentially different data management policy. As such, for every CPU/GPU/... target one has to define a new Workload object. They can of course share some parts, for example the payload IR etc.

[rengolin]

Summary of some of the discussion we had internally:

• The Workload class is a wrapper over other interface (payload, schedule, etc). The services it provides are to glue the interfaces that will be implemented by third-parties. It may be an interface, if we want to have different wrappers, but ultimately, that glue code will be very similar.
• The Payload interface represents the program being executed, in MLIR form. It will need extra information about the entry point and its arguments and return types, and how to handle data initialization and allocations. This will be different for CPUs, GPUs and other devices, but it could also be different across frameworks on the same kind of device, so it could be a separate interface that composes into the Workload class, or as an argument to the Payload class.
• The Schedule interface would share some of the common methods from Payload due to it being also MLIR, but the considerations are very different. We also discussed having a separate PassManager interface, and bundle them into a higher-level Pipeline interface, that is used by the optimizer and auto-tuner.
• The Execution interface would be responsible for collecting environment settings, required libraries and tools, and to execute the payload, optionally providing execution stats (for performance evaluation). That last item is hard to be precise, so we may need to wing it for the initial period and then coalesce into a design after we create a monster.
• The Benchmark infrastructure would probably use some off-the-shelf benchmarking runner helper and binding the Workload class to collect the inputs (payload, pipeline, dependencies), running the interface methods (compile, optimize, tune, run, etc), and collecting some stats in a pre-agreed shape.

Lighthouse would provide the initial interfaces and glue code, and implement a basic structure for generic CPU/GPU execution, then allow forks to override those interfaces. With time, we may consider moving some of those interfaces to MLIR proper, but I'd argue to keep the glue code in Lighthouse for the longer term.

[rengolin]

Also, relevant input:
https://github.com/llvm/lighthouse/wiki/Integrator