Add initial framework for TKW

[martin-luecke]

TODO: expand PR description, this is currently for enabling commenting

Detailed description of this project in shark_turbine/kernel/functional/README.md

[stellaraccident]

I saw this branch in my latest fetch and was intrigued. Looks interesting

[harsh-nod]

Thanks for pushing this. Just some minor comments but looks good overall.

[stellaraccident]

One naming nit: aside from some functional inspiration, the big idea here is a wave level programming model. If we're allocating the third letter of the prefix to what something is, probably tkw vs tkf. I doubt this will be the only thing that takes more of a functional view of things. And in general, that is a how, not a what.

[martin-luecke]

Thank you for the comments, they should now all be addressed.
I think the renaming @stellaraccident suggested makes sense so I renamed the project to tkw for now.

[martin-luecke]

I updated the IREE version here but the CI still seems to pick up a version which does not expose the amdgpu dialect. Do you know where to adjust this? @harsh-nod @stellaraccident

[stellaraccident]

Look to the failing CI run?

[martin-luecke]

Sorry, I rushed a bit to get to some other work.

I only now understood that the CI pulls nod-ai/SHARK and uses those requirements.
The CI here contains an sed command to adjust the SHARK requirements from iree-turbine/main to the currently tested commit.
The sed command is here:

iree-turbine/.github/workflows/test_shark.yml

Line 49 in 45f1de8

sed -i 's/iree-turbine#/iree-turbine.git@${{github.sha}}#/g' requirements.txt

and seems to be supposed to match here

However, it does not. I am not sure whether that is intended to test everything with the requirements of iree-turbine/main or an oversight.
I modified it to do the replacement, but that currently messes up torch dependencies. Also it still seems to not pull the iree version indicated in iree-requirements.txt. I assumed this is the place to modify the requirement, similar to how it is done in this patch
But I also see that in earlier CI runs unrelated to this PR a different iree version from the iree-requirements.txt is used for testing. e.g. here which pulls version 20240410.859 instead of version 20240514.893 indicated in the iree-requirements.

I will have to revisit this later.

[martin-luecke]

Thank you for the comments @harsh-nod
I will explore tomorrow autogenerating more of the operation boilerplate and in particular see whether the registration of the custom ops to the context can be done differently.

[martin-luecke]

Oops, seems my VSCode plugin glitched and I accidentally closed. Now reopened.
I pushed the lit tests to their own top-level directory as suggested to avoid test tools fighting each other.

[harsh-nod]

This is great Martin and apart from a few minor changes, seems good to land. Thanks for the push in getting this PR to this point. This will give us a great foundation to build on.

[martin-luecke]

Thank you @harsh-nod

Inspired by your comment I added the add_to_graph to our custom ops and took the time to simplify the interface for adding them to the graph after tracing.
In particular I removed the fx_op field from CustomOp which was required for initialization before. It is more a technical detail and remnant from the previous implementation. Instead, there is now a _tracing_function field on the class, which is set automatically by the define_op decorater.

With that the API for creating nodes after the tracing looks like this:

[...]
read = Read(a).add_to_graph(graph)
write = Write(read, a, 4).add_to_graph(graph)

[harsh-nod]

Thank you @harsh-nod

Inspired by your comment I added the add_to_graph to our custom ops and took the time to simplify the interface for adding them to the graph after tracing. In particular I removed the fx_op field from CustomOp which was required for initialization before. It is more a technical detail and remnant from the previous implementation. Instead, there is now a _tracing_function field on the class, which is set automatically by the define_op decorater.

With that the API for creating nodes after the tracing looks like this:

[...]
read = Read(a).add_to_graph(graph)
write = Write(read, a, 4).add_to_graph(graph)

Thanks Martin, that's awesome. Thank you for the push and this is a great first PR to build on top of.

[harsh-nod]

lgtm! thanks!