llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

torch: allow torch dialect ops after running drop-shape pass (#979) 

In the `pyhpc_turbulent_kinetic_energy` TorchBench benchmark, the shape
calculation occurs inside loops, but because `DropShapeCalculationsPass`
does not explicitly mark the Torch dialect as legal, the pass execution
fails.

This patch adds Torch to the list of legal dialects, and adds a test to
validate the translation.
pull/977/merge snapshot-20220626.515
Ashay Rane 2022-06-25 07:27:47 -07:00 committed by GitHub
parent 1be604bfd3
commit 163fa57cde
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 37 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lib/Dialect/Torch/Transforms/DropShapeCalculations.cpp
View File

@ -53,6 +53,7 @@ class DropShapeCalculationsPass
RewritePatternSet patterns(context);
patterns.insert<DropShapeCalculateOp>(context);
ConversionTarget target(*context);
target.addLegalDialect<Torch::TorchDialect>();
target.addIllegalOp<ShapeCalculateOp>();
target.addLegalOp<func::FuncOp>();

36
test/Dialect/Torch/drop-shape-calculations.mlir
View File

@ -19,3 +19,39 @@ func.func @basic(%arg0: !torch.vtensor<[2,?],unk>) -> !torch.vtensor {
%1 = torch.tensor_static_info_cast %0 : !torch.vtensor<[2,?],unk> to !torch.vtensor
return %1 : !torch.vtensor
}
// -----
// CHECK-LABEL: func.func @shape_calc_in_loop(
// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor<[2,?],unk>) -> !torch.vtensor<[2,?],unk> {
func.func @shape_calc_in_loop(%arg: !torch.vtensor<[2,?],unk>) -> !torch.vtensor<[2,?],unk> {
%one = torch.constant.int 1
// CHECK: %[[ONE:.*]] = torch.constant.int 1
%two = torch.constant.int 2
// CHECK: %[[TWO:.*]] = torch.constant.int 2
%true = torch.constant.bool true
// CHECK: %[[TRUE:.*]] = torch.constant.bool true
torch.prim.Loop %one, %true, init() {
// CHECK: torch.prim.Loop %[[ONE]], %[[TRUE]], init() {
^bb0(%in: !torch.int):
%shape_calc = torch.shape.calculate {
%tanh = torch.aten.tanh %arg : !torch.vtensor<[2,?],unk> -> !torch.vtensor<[2,?],unk>
torch.shape.calculate.yield %tanh : !torch.vtensor<[2,?],unk>
} shapes {
%size = torch.aten.size.int %arg, %one : !torch.vtensor<[2,?],unk>, !torch.int -> !torch.int
%list = torch.prim.ListConstruct %two, %size : (!torch.int, !torch.int) -> !torch.list<int>
torch.shape.calculate.yield.shapes %list : !torch.list<int>
} : !torch.vtensor<[2,?],unk>
// CHECK: torch.aten.tanh %[[ARG]] : !torch.vtensor<[2,?],unk> -> !torch.vtensor<[2,?],unk>
torch.prim.Loop.condition %true, iter()
// CHECK: torch.prim.Loop.condition %[[TRUE]], iter()
} : (!torch.int, !torch.bool) -> ()
return %arg : !torch.vtensor<[2,?],unk>
// CHECK: return %[[ARG]] : !torch.vtensor<[2,?],unk>
}