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/drop-shape-calculations.mlir

22 lines
1.2 KiB
MLIR
Raw Normal View History

Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
// RUN: torch-mlir-opt -torch-drop-shape-calculations -split-input-file %s | FileCheck %s
// CHECK-LABEL: func @basic(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor<[2,?],unk>) -> !torch.vtensor {
// CHECK: %[[TANH:.*]] = torch.aten.tanh %[[ARG]] : !torch.vtensor<[2,?],unk> -> !torch.vtensor<[2,?],unk>
// CHECK: %[[ERASED:.*]] = torch.tensor_static_info_cast %[[TANH]] : !torch.vtensor<[2,?],unk> to !torch.vtensor
// CHECK: return %[[ERASED]] : !torch.vtensor
func @basic(%arg0: !torch.vtensor<[2,?],unk>) -> !torch.vtensor {
%int2 = torch.constant.int 2
%int1 = torch.constant.int 1
%0 = torch.shape.calculate {
%2 = torch.aten.tanh %arg0 : !torch.vtensor<[2,?],unk> -> !torch.vtensor<[2,?],unk>
torch.shape.calculate.yield %2 : !torch.vtensor<[2,?],unk>
} shapes {
%2 = torch.aten.size.int %arg0, %int1 : !torch.vtensor<[2,?],unk>, !torch.int -> !torch.int
Elide `!torch.` prefix in nested dialect types. This leads to much more succinct types in many cases: ``` !torch.list<!torch.int> !torch.list<int> !torch.tuple<!torch.list<!torch.int>, !torch.list<!torch.int>> !torch.tuple<list<int>, list<int>> !torch.optional<!torch.list<!torch.int>> !torch.optional<list<int>> !torch.list<list<list<tensor>>> !torch.list<!torch.list<!torch.list<!torch.tensor>>> ``` I would like to take this further and allow omitting the `!torch.` prefix in all cases, but that's harder -- for example, we currently use `FuncOp` for functions, and so I don't think we can customize the printing there. It seems like it will be a longer road to getting that level of customization.
2022-03-16 07:22:56 +08:00
%3 = torch.prim.ListConstruct %int2, %2 : (!torch.int, !torch.int) -> !torch.list<int>
torch.shape.calculate.yield.shapes %3 : !torch.list<int>
Introduce new shape library design. See the documentation in `docs/shape_lib.md` and `docs/adding_a_shape_function.md` for an overview of the system. This completely overhauls how we represent shape functions. In particular, RefineTypes does not infer shapes anymore (only dtypes). Shape functions are now written in (TorchScript'able) Python. Recommended review order: 1. Read `docs/shape_lib.md` and `docs/adding_a_shape_function.md`. 1. Code and tests for ReifyShapeCalculations, DropShapeCalculations. 1. Code and tests for SimplifyShapeCalculations. 1. shape_lib_gen.py 1. Code and tests for new RefineTypes pass. 1. Random folders/canonicalizers in TorchOps.cpp and associated test in `canonicalize.mlir`. 1. New ReadOnly trait inferred from the registry. 1. Any miscellaneous remaining stuff. Example `-print-ir-after-all` for ElementwiseUnaryModule: [IR lowering dump](https://gist.github.com/silvasean/e4dc8cbc8d00aac7819602e3cbd8e212). Example `-print-ir-after-all` for ElementwiseBinaryModule: [IR lowering dump](https://gist.github.com/silvasean/daf6860ecced732af3568af6b1899113).
2022-03-10 08:44:22 +08:00
} : !torch.vtensor<[2,?],unk>
%1 = torch.tensor_static_info_cast %0 : !torch.vtensor<[2,?],unk> to !torch.vtensor
return %1 : !torch.vtensor
}