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/ATen/canonicalize.mlir

28 lines
1.2 KiB
MLIR
Raw Normal View History

Add recognition/folder/lowering for aten::__is__, aten::ne.int, and aten::dim Interestingly, TorchScript has its own op (`torch::jit::Operator`) registry separate from the dispatcher (it is a superset of the dispatcher). This is where the "prim" ops and some "aten" ops (that should probably be renamed to "prim") live. In particular, `aten::__is__` is in that latter category of "aten but really prim". This registry is also the source of truth for what the TorchScript interpreter calls into when it executes. The bulk of the "not part of the dispatcher" ops live in https://github.com/pytorch/pytorch/blob/09feb5f579d40b550559360870907963e20c6954/torch/csrc/jit/runtime/register_prim_ops.cpp#L82 And the registry itself lives in: https://github.com/pytorch/pytorch/blob/09feb5f579d40b550559360870907963e20c6954/torch/csrc/jit/runtime/operator.cpp#L196 This fold further reduces the IR of ResNet by folding away some more not-taken branches. These not-taken branches in ResNet require first-class handling of the list type which we don't yet have on any backend.
2021-04-28 06:15:50 +08:00
// RUN: npcomp-opt %s -canonicalize | FileCheck %s
// CHECK-LABEL: func @aten.__is__
// CHECK: %[[FALSE:.*]] = basicpy.bool_constant false
// CHECK: return %[[FALSE]] : !basicpy.BoolType
func @aten.__is__(%arg0: !basicpy.ListType, %arg1: !basicpy.NoneType) -> !basicpy.BoolType{
%0 = "aten.__is__"(%arg0, %arg1) : (!basicpy.ListType, !basicpy.NoneType) -> !basicpy.BoolType
return %0 : !basicpy.BoolType
}
Add aten::len.t, aten::size, and aten::gt.int primitive ops Also add some canonicalizations that finally reduce ResNet down to a single block.
2021-04-30 06:38:11 +08:00
// CHECK-LABEL: func @aten.size(
// CHECK: %[[CM1:.*]] = constant -1 : i64
// CHECK: %[[C3:.*]] = constant 3 : i64
// CHECK: %[[RET:.*]] = basicpy.build_list %[[CM1]], %[[C3]] : (i64, i64) -> !basicpy.ListType
// CHECK: return %[[RET]] : !basicpy.ListType
func @aten.size(%arg0: tensor<?x3xf32>) -> !basicpy.ListType {
%0 = "aten.size"(%arg0) : (tensor<?x3xf32>) -> !basicpy.ListType
return %0 : !basicpy.ListType
}
// CHECK-LABEL: func @aten.len.t(
// CHECK: %[[LENGTH:.*]] = constant 2 : i64
// CHECK: return %[[LENGTH]] : i64
func @aten.len.t(%arg0: i64) -> i64 {
%0 = basicpy.build_list %arg0, %arg0 : (i64, i64) -> !basicpy.ListType
%1 = "aten.len.t"(%0) : (!basicpy.ListType) -> i64
return %1 : i64
}