llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/frontends/pytorch/test/node_import/loop.py

73 lines
3.4 KiB
Python
Raw Normal View History

Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# -*- Python -*-
# This file is licensed under a pytorch-style license
# See frontends/pytorch/LICENSE for license information.
import torch
import torch_mlir
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
import typing
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# RUN: %PYTHON %s | npcomp-opt | FileCheck %s
mb = torch_mlir.ModuleBuilder()
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_Loop_forlike(
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK-SAME: %[[MAX_ITERATIONS:.*]]: !torch.int) -> f64 {
Add `!torch.bool` type. This finishes removing the dependence on the basicpy dialect! Changes: - Add `!torch.bool` type and replace use of `!basicpy.BoolType` in Torch-related code. - Rename BuiltinTensorize to BackendTypeConversion since now it handles bool conversions (and, when we add !torch.int and !torch.float, it will handle those as well), and generalize the related utilities (I also moved them to Torch/Transforms since they aren't really part of Torch/IR). - Add `torch.to_i1` and `torch.from_i1` ops for materializations - [cleanup] Reorganize `torch.constant.*` ops in TorchOps.td - Remove dependency of `torch` dialect on `basicpy` dialect and also `std` dialect. For `std`, we use some call related ops, but the `torch` dialect itself never produces them (we have passes that do though). This is fairly mechanical. Recommended review order: - New stuff in Torch/IR - New BuiltinTypeConversion files. - Mechnical fixups elsewhere.
2021-06-16 07:47:53 +08:00
# CHECK: %[[BOOL_TRUE:.*]] = torch.constant.bool true
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: %[[F_INIT:.*]] = torch.constant.float 0.000000e+00
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# CHECK: %[[RESULTS:.*]] = torch.prim.Loop %[[MAX_ITERATIONS]], %[[BOOL_TRUE]], init(%[[F_INIT]]) {
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: ^bb0(%[[IV:.*]]: !torch.int, %[[F_ITER:.*]]: f64):
# CHECK: %[[F_NEXT:.*]] = torch.aten.add.float_int %[[F_ITER]], %[[IV]] : f64, !torch.int -> f64
Miscellaneous changes while trying to work on ResNet18 - Move frontend lowering pipelines to c++ (this helps with reproducing failures in npcomp-opt) - Add debugging printouts when compilation fails on RefBackendTestConfig The experience now when a test fails during MLIR lowering is now like this: ``` NPCOMP TorchScript Object Graph IR -> NPCOMP Backend IR lowering failed with the following diagnostics: failed to legalize operation 'torch.global_slot' Module does not conform to npcomp's backend contract. See dialect conversion legality information above. Error can be reproduced with: $ npcomp-opt -torchscript-to-npcomp-backend-pipeline /tmp/ResNet18Module.mlir ``` And when TorchScript->MLIR import fails it looks like this: ``` PyTorch TorchScript module -> NPCOMP Object Graph IR import failed with the following diagnostics: unhandled prim operation: %18 : int = prim::min(%17) # /usr/local/google/home/silvasean/.local/lib/python3.9/site-packages/torch/nn/functional.py:4532:4 ``` Also, - Add `--filter=<regex>` to e2e test harness to filter tests. - Add a few prim ops that were needed to import ResNet18 - Fix torch.prim.Loop.condition assemblyFormat (it previously would not round-trip in the case of no loop-carried variables)
2021-04-22 06:07:15 +08:00
# CHECK: torch.prim.Loop.condition %[[BOOL_TRUE]], iter(%[[F_NEXT]] : f64)
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: } : (!torch.int, !torch.bool, f64) -> f64
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
# CHECK: return %[[RESULTS:.*]] : f64
@mb.import_function
@torch.jit.script
def prim_Loop_forlike(n: int):
f = 0.0
for i in range(n):
f += i
return f
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_Loop_whilelike(
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK-SAME: %[[VAL_0:.*]]: !torch.int) -> f64 {
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: %[[F_INIT:.*]] = torch.constant.float 3.200000e+00
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: %[[MAX_ITERATIONS:.*]] = torch.constant.int 9223372036854775807
# CHECK: %[[COND_INIT:.*]] = torch.aten.lt.float_int %[[F_INIT]], %[[VAL_0]] : f64, !torch.int -> !torch.bool
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK: %[[RET:.*]] = torch.prim.Loop %[[MAX_ITERATIONS]], %[[COND_INIT]], init(%[[F_INIT]]) {
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: ^bb0(%[[F_ITER:.*]]: !torch.int, %[[F_ITER:.*]]: f64):
Significantly restructure torch/aten import design. This is a really major and invasive restructuring of the way we get torch operators (`torch::jit::Operator` / `c10::OperatorHandle`) into MLIR. Please forgive the challenging review, but due to the sheer invasiveness, it wasn't really practical do do it in sane smaller pieces. This fully replaces everything that was already working on the TorchScript path (actually, more -- we added tanh support to TorchToLinalg in order to delete the older code paths). Additionally, I've kept the lights on for the acap path too, including what little e2e stuff was working before (for expediency I made a few tiny compromises along the way that will be easy to undo when we give that path proper attention). Overview of the new design: - The torch operator `somens::someunqualname.someoverloadname` is imported as `torch.somens.someunqualname.someoverloadname` (skip the last dotted part if the overload name is empty), OR, if we don't have such an op registered, it is imported as `torch.operator "somens.someunqualname.someoverloadname" (...) : ...`. - The addition of the "overload name" is a critical element here, as the `(ns,unqual,overload)` triple is unique, which solves a lot of problems we were having. - This involves having separate MLIR ops for the `trailing_` and `.out` variants and all the different overloads. This seemed necessary, because the set of overloads is so wild and varied and unstructured. The previous design was leaning into some underlying structure that just isn't there -- the default situation is the "random overload that we want to manage on the MLIR side", rather than that being an exception. E.g. `aten::ne` (not-equal) has 21 overloads, only 4 of which are c10 dispatcher ops see [gist](https://gist.github.com/silvasean/190ba918c550c956260e21254e1b8aa1), and the "out" variant is really called `.Tensor_out` instead of `.out` as it frequently is for other ops. - Rationale for all being in `torch` namespace: the set of operators are so varied and unstructured that "dialect per namespace" doesn't result in anything resembling the typical MLIR dialect boundary expectations. We could maybe draw the boundary at dispatcher ops vs non-dispatcher ops, but that doesn't seem to really result in very much useful structure at this point in time. - Note: within the torch operator registry, we effectively have a mini-basicpy subdialect (already type-resolved), which is reasonably structured. - The existing Torch op interfaces are also removed -- now that we track the overload name, we can losslessly find the original operator. - Instead of `ATenRecognizeKernelsPass`, we now have a `ReduceOpVariantsPass` that keys off certain traits (and perhaps eventually interfaces) to reduce variants of ops to a smaller set, ideally operating on immutable tensors and using surrounding ops to model the mutability/aliasing aspects. - Note: `torch.ns.unqual.overload` ops allow both immutable and mutable tensors (unlike the previous hard distinction in the common case). This is a premonition for a future change that will introduce a bona fide `!torch.tensor` type that will clean up a bunch of stuff. - `TorchToLinalg` / `TorchToStd` supercede the existing "ATen->TCF->TCP->Linalg" path. - The new `torch_ods_gen.py` supercedes `torch_signature_ods_gen.py`. It should look somewhat familiar, but the benefit of hindsight has allowed a lot of simplifications. The overall trend seems to be to make the `torch` dialect a nice layer independent of anything else. It feels like as a natural result of various future changes we will be removing the reliance on basicpy+numpy dialects and have a nice self-contained type system too that properly models the TorchScript type system (including proper subtyping, mutable/immutable tensors, optional dtype, etc.). Recommended review order: - Start at some of the new import IR, e.g. in `frontends/pytorch/test/node_import/prim.py`, `frontends/pytorch/test/acap_export/test_export_add3.py`, and other tests. - `frontends/pytorch/python/torch_mlir_utils/codegen/torch_ods_gen.py` and associated generated files: - `include/npcomp/Dialect/Torch/IR/GeneratedAtenOps.td` - `include/npcomp/Dialect/Torch/IR/GeneratedPrimOps.td` - Inspect `ReduceOpVariants.cpp` / `reduce-op-variants.mlir` and the new traits in `include/npcomp/Dialect/Torch/IR/TorchTraits.h` - Various code changes in the import path in `frontends/pytorch/csrc/builder`. Probably most interesting is the new code in `torch_to_mlir_utils.cpp` that has the logic to create the `torch.operator` ops or `torch.ns.unqual.overload` ops. This is the [new ResNet IR](https://gist.github.com/silvasean/5407aafb710d07612b7b5b92eabecebe), just to be able to look at a substantial sample of IR in the new style.
2021-05-05 05:42:50 +08:00
# CHECK: %[[F_NEXT:.*]] = torch.aten.mul.float %[[F_ITER]], %[[F_ITER]] : f64, f64 -> f64
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: %[[COND_ITER:.*]] = torch.aten.lt.float_int %[[F_NEXT]], %[[VAL_0]] : f64, !torch.int -> !torch.bool
Miscellaneous changes while trying to work on ResNet18 - Move frontend lowering pipelines to c++ (this helps with reproducing failures in npcomp-opt) - Add debugging printouts when compilation fails on RefBackendTestConfig The experience now when a test fails during MLIR lowering is now like this: ``` NPCOMP TorchScript Object Graph IR -> NPCOMP Backend IR lowering failed with the following diagnostics: failed to legalize operation 'torch.global_slot' Module does not conform to npcomp's backend contract. See dialect conversion legality information above. Error can be reproduced with: $ npcomp-opt -torchscript-to-npcomp-backend-pipeline /tmp/ResNet18Module.mlir ``` And when TorchScript->MLIR import fails it looks like this: ``` PyTorch TorchScript module -> NPCOMP Object Graph IR import failed with the following diagnostics: unhandled prim operation: %18 : int = prim::min(%17) # /usr/local/google/home/silvasean/.local/lib/python3.9/site-packages/torch/nn/functional.py:4532:4 ``` Also, - Add `--filter=<regex>` to e2e test harness to filter tests. - Add a few prim ops that were needed to import ResNet18 - Fix torch.prim.Loop.condition assemblyFormat (it previously would not round-trip in the case of no loop-carried variables)
2021-04-22 06:07:15 +08:00
# CHECK: torch.prim.Loop.condition %[[COND_ITER]], iter(%[[F_NEXT]] : f64)
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: } : (!torch.int, !torch.bool, f64) -> f64
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK: return %[[RET:.*]] : f64
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
@mb.import_function
@torch.jit.script
def prim_Loop_whilelike(n: int):
f = 3.2
while f < n:
f = f * f
return f
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK-LABEL: func @__torch__.prim_Loop_derefine(
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK-SAME: %[[ARG:.*]]: !torch.int) -> !torch.optional<!torch.int> {
Add `!torch.bool` type. This finishes removing the dependence on the basicpy dialect! Changes: - Add `!torch.bool` type and replace use of `!basicpy.BoolType` in Torch-related code. - Rename BuiltinTensorize to BackendTypeConversion since now it handles bool conversions (and, when we add !torch.int and !torch.float, it will handle those as well), and generalize the related utilities (I also moved them to Torch/Transforms since they aren't really part of Torch/IR). - Add `torch.to_i1` and `torch.from_i1` ops for materializations - [cleanup] Reorganize `torch.constant.*` ops in TorchOps.td - Remove dependency of `torch` dialect on `basicpy` dialect and also `std` dialect. For `std`, we use some call related ops, but the `torch` dialect itself never produces them (we have passes that do though). This is fairly mechanical. Recommended review order: - New stuff in Torch/IR - New BuiltinTypeConversion files. - Mechnical fixups elsewhere.
2021-06-16 07:47:53 +08:00
# CHECK: %[[TRUE:.*]] = torch.constant.bool true
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: %[[NONE:.*]] = torch.constant.none
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: %[[NONE_DEREFINED:.*]] = torch.derefine %[[NONE]] : !torch.none to !torch.optional<!torch.int>
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
# CHECK: %[[RET:.*]] = torch.prim.Loop %[[ARG]], %[[TRUE]], init(%[[NONE_DEREFINED]]) {
Add `!torch.int` type. This replaces the ad-hoc use of `i64` throughout the Torch layer, and helps to keep it crystal clear the distinction between `!torch.int` (which is modeling the Python `int` type) and the various types that serve as dtypes of tensors, which are a totally different type universe. Changes: - `!torch.int` type and C bindings. - Change `torch.constant.int` parser to not need the `: i64` at the end. - `m_TorchConstantInt` matcher to aid with matching constants. - BackendTypeConversion changes for `!torch.int` -> `i64` type conversion. - Refactor finalizing patterns in FinalizingBackendTypeConversionPass (they were getting very repetitive). - Mechanical rewriting of `!torch.int` to `i64` in all the tests, and `AnyTorchIntType` to `Torch_IntType` in the `.td` files.
2021-06-17 06:53:15 +08:00
# CHECK: ^bb0(%[[IV:.*]]: !torch.int, %[[X_ITER:.*]]: !torch.optional<!torch.int>):
# CHECK: %[[X_NEXT:.*]] = torch.derefine %[[ARG]] : !torch.int to !torch.optional<!torch.int>
# CHECK: torch.prim.Loop.condition %[[TRUE]], iter(%[[X_NEXT]] : !torch.optional<!torch.int>)
# CHECK: } : (!torch.int, !torch.bool, !torch.optional<!torch.int>) -> !torch.optional<!torch.int>
# CHECK: return %[[RET:.*]] : !torch.optional<!torch.int>
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
@mb.import_function
@torch.jit.script
def prim_Loop_derefine(n: int):
x: typing.Optional[int] = None
for i in range(n):
x = n
return x
Add support for prim::Loop op. This is a funny one. It combines a `for` and `while` loop in one op. We will need to write some conversions to `scf`.
2021-03-02 07:00:32 +08:00
mb.module.operation.print()
print()