Propagate device data names

[antoniojkim]

Need to propagate device data names in order to be able to distinguish between inputs and weights.

An example of its intended usage has been added to the reference backend. The idea is to assign names to tensors on the python side (e.g. perhaps in the dataloader) so that you know which tensors are which in the actual compile

Also, export remaining headers that weren't exported before.

CC: @ke1337 @henrytwo

[silvasean]

can you add a test for this?

[silvasean]

I'm not sure I understand what the test is testing. It seems to call the API but I don't see it checking the result.

[antoniojkim]

I'm not sure I understand what the test is testing. It seems to call the API but I don't see it checking the result.

There's unfortunately no good way to check the result. I originally checked that the number of inputs was greater than zero, but this is not true for all ops apparently and thus some ops fail that check.

I also tried saving the input names and checking them against the values that are set, but this also doesn't work as multiple tests are run in parallel and thus there is a race condition in the storage and retrieval of this value.

As such, I decided to only add it as a spot check. Especially considering that this is not something that is needed for all lazy backends (for example, the TorchScript backend has no use for this). I figured its fine to just print the number of inputs in the reference lazy backend so that we can do a visible check of the output to verify that its working as intended.

@silvasean please let me know if my reasoning for this is sufficient

[antoniojkim]

@silvasean Gentle reminder to please review again

[silvasean]

Sorry, I've been OoO.

What I meant to say is that I see std::cout << num_inputs << " input tensors found" << std::endl; but I don't see where you check that that gets actually printed. And I don't think that calling this API in the end-to-end tests is very useful.

I would recommend that you add unit tests like we did with the "eager mode" here: https://github.com/llvm/torch-mlir/tree/main/python/test/eager_mode

I'm imagining a 3 line test where you call lazy_backend.set_parameter_name(e, f"some_name") and then either print the graph or print the node or get the parameter name or something and print it out and check that "some_name" is in the output.

[antoniojkim]

What I meant to say is that I see std::cout << num_inputs << " input tensors found" << std::endl; but I don't see where you check that that gets actually printed. And I don't think that calling this API in the end-to-end tests is very useful.

There is no concrete check. I only meant for it to be a "check by inspection". Its not required by all backends and I thought it was sufficient to be able to see it in the test logs.

I would recommend that you add unit tests like we did with the "eager mode" here: https://github.com/llvm/torch-mlir/tree/main/python/test/eager_mode

The problem with doing this is that the tests are run in parallel in the same process. This causes race conditions as there is only a single instance of the static lazy tensor C++ symbols.

I'm imagining a 3 line test where you call lazy_backend.set_parameter_name(e, f"some_name") and then either print the graph or print the node or get the parameter name or something and print it out and check that "some_name" is in the output.

This is what I was hoping to add as well, but with the current test infrastructure, unless I'm mistaken, I don't think this is possible due to the reasons I listed above.

[antoniojkim]

@silvasean I've added some unit tests. Please review again when you get a chance