RFC: Upcoming deprecation of iree-run-trace/iree-benchmark-trace

[benvanik]

Many moons ago we built out a way to simulate pipelines of stateful invocations for testing via the iree-run-trace/iree-benchmark-trace tools. These take a YAML file with the expected inputs, optionally outputs (though we don't do anything with them today), and has some basic support for blackboard-style state.

Since it was originally written we've gotten multiple modules working much better and it's now possible to do nearly everything the YAML could do natively just using modules: https://github.com/openxla/iree/tree/main/samples/multiple_modules
By using modules state, loops, integration with other tooling (check tests), and all features are fully supported (parameters, multiple devices, async, collectives, etc). Anything that can run IREE programs can run the pipelines meaning that iree-run-module/iree-benchmark-module/iree-check-module, the python API, etc can consume them without changes.

As an example, the following YAML trace:
https://github.com/openxla/iree/blob/main/tools/test/iree-run-trace.yml
can be written as:
```
func.func private @module.mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
func.func @run(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst0 = arith.constant dense<[0,1,2,3]> : tensor<4xf32>
  %0 = call @module.mul(%arg0, %cst0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  %cst1 = arith.constant dense<3> : tensor<4xf32>
  %1 = call @module.mul(%0, %cst1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  return %1 : tensor<4xf32>
}
```
the pipeline is then compiled and run by passing both the module under test and the pipeline module:
```
iree-compile pipeline.mlir -o=pipeline.vmfb
iree-run-module --module=module.vmfb --module=pipeline.vmfb --function=run
```

We never really got much tooling support around the YAML traces but the idea was that bindings could write out the traces as a recording of the call sequence by the user with just stdout/printf. The idea now is that just as a binding layer could write out textual YAML lines it could also write out textual MLIR statements - it's just a different syntax for the same thing but with verification and the full fidelity of the system. In other words, there's no need to pull in MLIR IR builders and such and instead just emit them textually.

The loose plan is to deprecate the YAML trace replay tooling by removing the top level tools (iree-run-trace/iree-benchmark-trace) from deployments and the repo. The end-to-end matmul tests currently use the same infrastructure and we'll have to move them over to being modules but to start the code will be moved to a test-only path for eventual reworking. Once reworked we can entirely remove the libyaml dep from the repo.

Since it's pretty tricky to use and doesn't have optimized workflows we don't believe anyone is using this functionality (outside of the in-tree matmul tests) but wanted to give a heads up in case someone is to ensure we carry forward any features people are relying on.

[Scott Todd]

This generally SGTM. There was an old issue tracking documenting those tools and smoothing out some of their rough edges that never got much attention: https://github.com/openxla/iree/issues/6688

One of the nice selling points of the YAML traces was that they reflected the execution flow directly and could be used to reproduce subtle VM runtime issues deterministically. I'm wondering if textual MLIR statements with the same purpose could drift from the actual sequence of calls, either intentionally or accidently (e.g. via MLIR inlining / canonicalization / IREE const-eval). I guess on the other side of that, maybe we could use the mlir-reduce style tools on "MLIR trace" files where we couldn't use them on YAML files before.

[Jacques Pienaar]

> We never really got much tooling support around the YAML traces but the idea was that bindings could write out the traces as a recording of the call sequence by the user with just stdout/printf

So this doesn't actually exist yet with the tracing tool and so it was manual authored before and manual authored after, but in a way that would provide a more flexible execution model? (and one potentially not limited to just testing?).

--
You received this message because you are subscribed to the Google Groups "iree-discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to iree-discuss...@googlegroups.com.
To view this discussion on the web visit https://groups.google.com/d/msgid/iree-discuss/f94dab83-e3d1-47fa-9659-462fdc3ac0dcn%40googlegroups.com.

[Jacques Pienaar]

Was just checking, would this then also result in deprecating https://github.com/openxla/iree/blame/c859e290c649b7fdda897b23cb5b99a26cd1053c/tests/e2e/matmul/BUILD.bazel#L28 ? (was just trying to trace the tool & yaml usage)

[Ben Vanik]

When compiling multiple modules we don't do any inlining or anything - so making a call across module boundaries should be (effectively) identical to any other way you could make a call (like the trace replay infra using the iree_vm_invoke).

I've ported over the e2e matmul test infra to emit MLIR instead of the YAML traces; now the python generator emits something like this which we compile and run:

https://gist.github.com/benvanik/8e637b2e10241d54ef5d4a6e567e41fa

the tool implements a custom module with generate_random_matrix/check_matmul_results that do what the magic YAML pseudorandom number generator helper did and the existing e2e matmul test runner tool did only now with proper support for non-host-visible buffers (this has been blocking some allocator improvements)

Branch PR here that fully eliminates our libyaml dep: https://github.com/openxla/iree/pull/16134

[Jacques Pienaar]

Great, I'd be happy to drop that dep :)

To view this discussion on the web visit https://groups.google.com/d/msgid/iree-discuss/d7c51836-e5ff-40b4-87ee-28b6adc0991en%40googlegroups.com.