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/torch-function-to-torch-bac...

18 lines
906 B
MLIR
Raw Normal View History

update llvm tag to a2620e00. (#1567) - also update MHLO to 57ba12a2(branch greencommit/2022-11-07-a2620e00) - change -pass-pipeline format to make tests pass.
2022-11-10 18:39:28 +08:00
// RUN: torch-mlir-opt -pass-pipeline='builtin.module(torch-function-to-torch-backend-pipeline{backend-legal-ops=torch.aten.square,torch.aten.argmax})' -split-input-file %s | FileCheck %s
Add a way for backends to control which ops are legal for them. We were already hitting many cases where backends different in terms of the legal ops that they wanted. This caused unnecessary coupling between the backends. Examples: - https://github.com/llvm/torch-mlir/pull/1161 - https://github.com/llvm/torch-mlir/pull/862 This PR centralizes all compilation to go through `torch_mlir.compile` so that we can keep the logic centralized there. We should move these lists closer to each backend. Especially cases like https://github.com/llvm/torch-mlir/pull/862 where blocking a decomposition is necessary to avoid a crash emphasize that the set of decompositions is tightly coupled to the backend, and should be "controlled by the backend" and not something arbitrarily tweakable. Also: - Fix a small bug in the way we passed through the backendLegalOps option. - Add better error messages in `torch_mlir.compile` for import errors.
2022-08-19 08:01:54 +08:00
// CHECK-LABEL: func.func @torch.aten.square
func.func @torch.aten.square(%arg0: !torch.vtensor<[?,?,?],f32>) -> !torch.vtensor<[?,?,?],f32> {
// CHECK: torch.aten.square
%0 = torch.aten.square %arg0 : !torch.vtensor<[?,?,?],f32> -> !torch.vtensor<[?,?,?],f32>
return %0 : !torch.vtensor<[?,?,?],f32>
}
// CHECK-LABEL: func.func @torch.aten.argmax
func.func @torch.aten.argmax(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor<[1,?],si64> {
%int0 = torch.constant.int 0
%true = torch.constant.bool true
// CHECK: torch.aten.argmax
%0 = torch.aten.argmax %arg0, %int0, %true : !torch.vtensor<[?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[1,?],si64>
return %0 : !torch.vtensor<[1,?],si64>
}