StableHLO support

[abeleinin]

Hi! I saw you added #85 yesterday, and I think StableHLO should be a lot simpler to add before we tackle ONNX. This initial pr just adds support for basic arithmetic ops, but I think the structure is here to easily extend it. I'm not sure how complete you want this crate to be before it lands, but I'll continue to add ops in the coming days.

I did have some issues with the lifetime of GraphTensors when constructing the Graph inside import_hlo. I'm not sure if this is from my lack of rust knowledge, but I did notice GraphTensor uses unsafe mutable references. All else equal when I constructed the graph in the function I was getting Tensor not found in the graph! when I tried to print inputs["0"] in the tests.

Also, I decided to return a HashMap so you can easily set function args by their names. I'm open to changing it to a Vec if that's your preference.

[jafioti]

@abeleinin This is awesome!

This is a really great start. I think we can take a look at some popular models represented in StableHLO and see what ops they use, if we can support the ops in at least a couple non-trivial models, I think it's good enough to merge. I don't know too much about StableHLO so I don't know how many ops this is, but my guess is that it's mostly a 80/20 situation where we can support most of the popular models with only the most used ops, and then add the rest in later.

For the GraphTensor thing, basically they contain an unsafe pointer back to the graph which lets us build the graph using graph tensors, and are only supposed to live as long as the Graph lives. If you're moving the Graph around after creation, it's probably a good idea to Box it using Box::new(graph) which will heap allocate it.

For the Tensor not found in graph error, that's happening if you're trying to access data for that tensor without it ever being set. Data is either set for a tensor manually with tensor.set(...) or during graph execution. For the tensor to keep its data after graph.execute() the tensor needs to be marked as retrieve by using tensor.retrieve(), which will allow you to do tensor.data() after the graph has ran. We should definitely clarify this better in the docs because I know this differs from a lot of other libs.

One other thing, I don't think you need to include luminal_hlo in the dependencies for luminal_hlo, because that just makes it a dependency on itself!

[abeleinin]

Cool! Thanks for the info. I was looking through the examples to figure out what would be good to support, maybe MobileNets or BERT would be good baselines.

[jafioti]

Yeah those make sense as solid example models

[abeleinin]

Sorry for the month long hiatus, getting back to work on this!

I just uploaded a README that outlines how I've been exporting models from PyTorch, which is a lot more painful than I expected. There's some nuance to the mlir structure that's a bit annoying, so you have to flatten it using the stablehlo-opt executable from the openxla/stablehlo github repo. Also, I include a simple squeezenet1_0.mlir module in this pr to have some more guidance on initial support. We only need to support the following StableHLO ops to get this working:

broadcast_in_dim, concatenate, constant, convolution, maximum, reduce, reduce_window, reshape

[abeleinin]

End-to-end squeezenet1_0 support in luminal via stablehlo mlir!

cargo run --example squeezenet examples/data/hamster.jpg

Output:

Top-5:
  #1: hamster (idx: 333, prob: 0.8738)
  #2: porcupine (idx: 334, prob: 0.0582)
  #3: Norwich Terrier (idx: 186, prob: 0.0173)
  #4: Yorkshire Terrier (idx: 187, prob: 0.0114)
  #5: guinea pig (idx: 338, prob: 0.0113)

Working on cleaning this up and improving the documentation soon. Also, I plan to look into what ops (if any) are missing in more advanced models (BERT, MobileNets, etc).