Support pytorch extensions

[rdadolf]

PyTorch allows new operators to be registered dynamically in modules. Torch-mlir already makes it fairly straightforward to add support for new operators, and this commit just extends that support to allow new PyTorch ops to come from an external module.

Previously, there was no hook to allow extensions to be loaded before querying the op registry, so extension ops would never be generated. Moreover, a bug caused ops imported from an extension to be impossible to support via shape inference.

This does not allow ops to be dynamically loaded into torch-mlir. Torch-mlir must still be compiled with support built-in.

This should make tasks like this one in #462 (and to a lesser extent #504) somewhat easier.

[rdadolf]

@silvasean Since you originally wrote most of these pieces, I'd appreciated an opinion.

[powderluv]

drive-by thank you for the PR :D

[silvasean]

This looks awesome. Before committing this, I would like to somehow make it load-bearing for the project so it gets tested/etc. To do that, I think we should add a minimal custom op end-to-end, which I think would consist of:

1. build Python module defining custom torch op
2. add stuff in ods_gen and shape_lib_gen for that op
3. dtype refinement in RefineTypes
4. TorchToLinalg/CustomOpsExample.cpp

I think it could be as simple as a custom unary op -- something for folks to follow the bread crumbs when adding their own. I think the QPyTorch folks are going to want to follow this too.

[rdadolf]

I wrote and tested all of that (in fact, it's literally written for a unary no-op---seems we had the same idea), but two things stopped me from including it in the PR:

1. I had no idea where to put it for testing purposes.
2. It means the test op will be exposed to the "real" world, which seemed unclean at the time.

If you folks don't mind, I don't mind.

Edit: this was my working branch that includes all of the non-torch-extension pieces of the code. I extracted this PR from that branch.

[silvasean]

• I had no idea where to put it for testing purposes.

What form does the custom op take? Is it just a Python module? Or is there C++ code involved?

• It means the test op will be exposed to the "real" world, which seemed unclean at the time.

Yeah, I'm sympathetic to this. However, PyTorch's op list is already pretty "unclean" and has all sorts of random stuff with varying levels of "core-ness" vs "custom-ness". I don't think that torch_mlir_custom.nop is that much more unclean than aten._foo_cudnn_special_thing

[rdadolf]

Sounds good. I'll start adding the rest of the code into the PR.

What form does the custom op take? Is it just a Python module? Or is there C++ code involved?

My example was a python wrapper around a C++ extension, similar to the tutorial from the PyTorch docs. It's small, but it does include a cmake file, requires compilation, etc.

The short version is:

Python side:

import os
import torch

# Register custom.nop as a side-effect of importing this module.
current_dir = os.path.dirname(os.path.abspath(__file__))
lib = os.path.join(*[current_dir, os.pardir, os.pardir, 'custom_nop', 'libnop.so'])
torch.ops.load_library(lib)

C++ side:

#include <torch/script.h> // One-stop header.

#include <iostream>

torch::Tensor nop(torch::Tensor t) {
  std::cout << "no-op executed\n";
  return t;
}

TORCH_LIBRARY(custom, m) {
  m.def("nop(Tensor t) -> Tensor");
  m.impl("nop", &nop);
}

[silvasean]

Ah, ok. Then I think this can be a standalone build inside build_tools. Perhaps build_tools/custom_op_example. And for clarity, let's make sure that the "namespace" of the op is always torch_mlir_custom_op_example, e.g. torch_mlir_custom_op_example.nop.

Or for simplicity, if there is a way to add this as a subdirectory of the existing build, perhaps bside the jit ir importer, that would be convenient and avoid a manual step to build build_tools/custom_op_example.

[rdadolf]

All set for merge, I think. Happy to take any last comments or requests.

[silvasean]

Awesome, thanks Bob!!

Can you squash the history before merging to main?

[rdadolf]

Can you squash the history before merging to main?

Done. I don't have write access to the repo, so I'll need someone to hit the button. Thanks!

[silvasean]

Can you squash the history before merging to main?

Done. I don't have write access to the repo, so I'll need someone to hit the button. Thanks!

Just invited you. Can't believe you didn't have it!