llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/test/Dialect/TorchConversion/annotate-abi.mlir

49 lines
1.6 KiB
MLIR
Raw Normal View History

Add basic support for lists. This plumbs through a vertical slice of support for lists. The main chunk of new code here is AnnotateABIPass which captures the program signature at the Torch backend contract layer, right before we start `TorchConversion`. The `TorchConversion` lowering process is lossy w.r.t. types, so it's necessary to do this for all targets in general. Like using `!iree.list` directly, we use IREE's ABI annotation representation for this, although there is nothing very IREE-specific about it (see https://github.com/google/iree/blob/main/docs/developers/design_docs/function_abi.md) We change `ListLiteralModule_basic` to use `!torch.int` because IREE doesn't support f64 yet (and we don't yet have a way for users to say that they want `!torch.float` to lower as f32). Recommended review order: - AnnotateABIPass and tests - Arg marshaling in npcomp_backend.py and `iree.py` - Updates to `list_programs.py` / `xfail_sets.py` - Moving DeleteDeadIREEListsPass to Backend/Common, so that backends that don't support lists can use it. RefBackend uses that pass, for example.
2021-09-03 03:47:59 +08:00
// RUN: npcomp-opt -split-input-file -verify-diagnostics %s -torch-annotate-abi
// -----
// CHECK-LABEL: builtin.func @basic_arg_and_ret(
// CHECK-SAME: attributes {iree.abi = "{\22a\22:[\22f64\22,\22i64\22,\22i1\22],\22r\22:[\22f64\22,\22i64\22,\22i1\22],\22v\22:1}"} {
builtin.func @basic_arg_and_ret(%arg0: !torch.float, %arg1: !torch.int, %arg2: !torch.bool) -> (!torch.float, !torch.int, !torch.bool) {
return %arg0, %arg1, %arg2 : !torch.float, !torch.int, !torch.bool
}
// -----
// CHECK-LABEL: builtin.func @list(
// CHECK-SAME: attributes {iree.abi = "{\22a\22:{{\[\[}}\22py_uniform_list\22,\22f64\22]],\22r\22:[],\22v\22:1}"} {
builtin.func @list(%arg0: !torch.list<!torch.float>) {
return
}
// -----
// CHECK-LABEL: builtin.func @tuple(
// CHECK-SAME: attributes {iree.abi = "{\22a\22:{{\[\[}}\22pytuple\22,\22f64\22,\22i64\22]],\22r\22:[],\22v\22:1}"} {
builtin.func @tuple(%arg0: !torch.tuple<!torch.float, !torch.int>) {
return
}
// -----
// CHECK-LABEL: builtin.func @dict(
// CHECK-SAME: attributes {iree.abi = "{\22a\22:{{\[\[}}\22py_uniform_dict\22,\22pystr\22,\22f64\22]],\22r\22:[],\22v\22:1}"} {
builtin.func @dict(%arg0: !torch.dict<!torch.str, !torch.float>) {
return
}
// -----
// CHECK-LABEL: builtin.func @tensor(
// CHECK-SAME: attributes {iree.abi = "{\22a\22:{{\[\[}}\22ndarray\22,\22f32\22,2,null,3]],\22r\22:[],\22v\22:1}"} {
builtin.func @tensor(%arg0: !torch.tensor<[?,3],f32>) {
return
}
// -----
// expected-error @+1 {{at function argument 0: unimplemented: ABI annotation for type '!torch.any'}}
builtin.func @unsupported(%arg0: !torch.any) {
return
}