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/Torch/GlobalizeObjectGraph/module-uses.mlir

49 lines
2.1 KiB
MLIR
Raw Normal View History

Extend GlobalizeObjectGraph to handle torch.prim.GetAttr returning NnModuleType This happens in practice. With this, we can globalize slots for the non-trivial classifier layer obtained from https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb This also adds support for tuple return types, which were needed by that model.
2021-02-19 09:10:17 +08:00
// RUN: npcomp-opt -torch-globalize-object-graph -split-input-file %s | FileCheck %s
torch.class_type @child {
torch.attr "float" : f64
}
torch.class_type @parent {
torch.attr "m" : !torch.nn.Module<"child">
torch.method "get_attr_returns_module_type", @get_attr_returns_module_type
torch.method "module_type_argument", @module_type_argument
torch.method "method_call", @method_call
}
// CHECK-LABEL: func @get_attr_returns_module_type() -> f64 {
func private @get_attr_returns_module_type(%arg0: !torch.nn.Module<"parent">) -> f64 {
%0 = torch.prim.GetAttr %arg0["m"] : !torch.nn.Module<"parent"> -> !torch.nn.Module<"child">
// CHECK-NEXT: %[[V:.*]] = torch.global_slot.get @m.float : f64
%1 = torch.prim.GetAttr %0["float"] : !torch.nn.Module<"child"> -> f64
// CHECK-NEXT: torch.global_slot.set @m.float = %[[V]] : f64
torch.prim.SetAttr %0["float"] = %1 : !torch.nn.Module<"child">, f64
// CHECK-NEXT: return %[[V]] : f64
return %1 : f64
}
// CHECK-LABEL: func @module_type_argument(
Introduce native `!torch.none` type. - Add `torch.constant.none` op to construct it (naming is chosen to be analogous to Torch's representation of a prim::Constant with NoneType, rather than using the "singleton" terminology of Basicpy).
2021-06-15 02:36:10 +08:00
// CHECK-SAME: %[[F:.*]]: f64) -> !torch.none {
func private @module_type_argument(%arg0: !torch.nn.Module<"parent">, %arg1: !torch.nn.Module<"parent">, %arg2: f64, %arg3: !torch.nn.Module<"parent">) -> !torch.none {
%0 = torch.constant.none
return %0 : !torch.none
Extend GlobalizeObjectGraph to handle torch.prim.GetAttr returning NnModuleType This happens in practice. With this, we can globalize slots for the non-trivial classifier layer obtained from https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb This also adds support for tuple return types, which were needed by that model.
2021-02-19 09:10:17 +08:00
}
Introduce native `!torch.none` type. - Add `torch.constant.none` op to construct it (naming is chosen to be analogous to Torch's representation of a prim::Constant with NoneType, rather than using the "singleton" terminology of Basicpy).
2021-06-15 02:36:10 +08:00
// CHECK-LABEL: func @method_call() -> !torch.none {
func private @method_call(%arg0: !torch.nn.Module<"parent">) -> !torch.none {
Add `torch.constant.int` and `torch.constant.float`. - This removes reliance on basicpy.numeric_constant. - Also, add OpAsmOpInterface to the `torch.constant.none` and `torch.constant.str` ops.
2021-06-16 03:42:51 +08:00
// CHECK-NEXT: %[[C:.*]] = torch.constant.float 4.300000e+01
%c = torch.constant.float 43.0
Introduce native `!torch.none` type. - Add `torch.constant.none` op to construct it (naming is chosen to be analogous to Torch's representation of a prim::Constant with NoneType, rather than using the "singleton" terminology of Basicpy).
2021-06-15 02:36:10 +08:00
// CHECK-NEXT: %[[F:.*]] = call @module_type_argument(%[[C]]) : (f64) -> !torch.none
%0 = call @module_type_argument(%arg0, %arg0, %c, %arg0) : (!torch.nn.Module<"parent">, !torch.nn.Module<"parent">, f64, !torch.nn.Module<"parent">) -> (!torch.none)
// CHECK-NEXT: return %[[F]] : !torch.none
return %0 : !torch.none
Extend GlobalizeObjectGraph to handle torch.prim.GetAttr returning NnModuleType This happens in practice. With this, we can globalize slots for the non-trivial classifier layer obtained from https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb This also adds support for tuple return types, which were needed by that model.
2021-02-19 09:10:17 +08:00
}
Add `torch.constant.int` and `torch.constant.float`. - This removes reliance on basicpy.numeric_constant. - Also, add OpAsmOpInterface to the `torch.constant.none` and `torch.constant.str` ops.
2021-06-16 03:42:51 +08:00
%c42 = torch.constant.float 42.0
Extend GlobalizeObjectGraph to handle torch.prim.GetAttr returning NnModuleType This happens in practice. With this, we can globalize slots for the non-trivial classifier layer obtained from https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb This also adds support for tuple return types, which were needed by that model.
2021-02-19 09:10:17 +08:00
%child = torch.nn_module {
torch.slot "float", %c42 : f64
} : !torch.nn.Module<"child">
torch.nn_module {
torch.slot "m", %child : !torch.nn.Module<"child">
} : !torch.nn.Module<"parent">