Binary payload in MLIR files

MLIR
1

@River707 @clattner

How is the binary payload in MLIR proposal going?

I have the following use-case that I think really needs that feature to not resort to horrible workarounds.

I want to compare the output of my MLIR compiler with a hand-coded C++ implementation of the same kernels, using libxsmm and known to be the optimal solution to each of those problems.

So I take MLIR as linalg-on-tensors, run my passes, and execute on the CPU using a modified version of mlir-cpu-runner that wraps the kernel into a benchmark loop. Our passes should basically repeat the same blocking, tiling and fusing as the ninja-code does, and so the timings must be the same, or very close.

But we also need to validate the outputs, and to catch corner cases on the blocking/tiling algorithms, we need random inputs, and a way to extract the outputs with enough precision. Right now, the inputs are hard-coded and the outputs are printed as text, which is not meaningful in any way.

We’re looking into using the same random seeds for the inputs (and same generator), but still no solution for the outputs.

An easy way out would be to create the tensors as binary files by a script, then run the C++ benchmark mmaping them. On the MLIR side, however, we don’t have that option. If we had a way to dump binary files into and out of MLIR files, we could use the same mechanism, and compare binary files directly (as floating-point values ± epsilon, etc).

However, this also means two things: first, the MLIR file has to be writable, and the JIT engine has to have access to the original file. I imagine this has to be true for training on MLIR files, but I’m not sure that’s part of the proposal.

Is this a reasonable use-case for the feature you guys are planning with the binary payloads in MLIR?

Any advice for extracting binary output from MLIR execution (without having to write my own runtime library)?

2

Have you looked at Resources? ([mlir] Add a new builtin DenseResourceElementsAttr · llvm/llvm-project@995ab92 · GitHub added nice helper too, and they are mmap’d if you use the bytecode format [mlir:Bytecode] Add shared_ptr<SourceMgr> overloads to allow safe mma… · llvm/llvm-project@18546ff · GitHub). Resources can also be unmanaged so in your runner loop you can change the value inside the Resource without changing the MLIR file. At least for input path, output path your runner could dump the binary output format corresponding to the resource by just invoking the bytecode … or even simpler if your program has the resource of interest, just mutate in memory and write out again (you’ll see in the mmap change that allowing mutation was considered but not yet implemented).

3

I had not seen this, looks like what I’m looking for. Thanks!

However, the test on the first commit is quite minimal, and the second commit has no tests at all. It’s hard to see how this could be used from IR or a builder.

Are there any examples (of IR or code) that use this feature?

4

Agreed we don’t have too many upstream/in OSS (not to imply its different internal my side, but internal elsewhere). From IR, mlir/test/IR/dense-resource-elements-attr.mlir is the only one (well mlir/test/CAPI/ir.c also shows format in textual format, running these via mlir-opt -emit-bytecode to get other format), in code mlir/unittests/IR/AttributeTest.cpp is the only one. A good example would be a good contribution :slight_smile: (I was going to make one that does training loop to show it all self-contained, but that’s currently ~7 from top of stack and so may not happen soon).

1 Like
5

All the logic and support is upstream, but River is the best to provide guidance here. He is out on vacation still but is back next week.

-Chris