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

[MLIR][TORCH] Add support for dynamic shape for Onnx.Transpose op (#2803) 

Signed-Off By: Vivek Khandelwal <vivekkhandelwal1424@gmail.com>
pull/2813/head
Vivek Khandelwal 2024-01-26 23:16:54 +05:30 committed by GitHub
parent 4964977e85
commit da7c6d2c16
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 18 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
View File

@ -1244,6 +1244,12 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
current[i] = i;
}
// Convert dynamic shape dimension.
for (unsigned i = 0; i < shape.size(); i++){
if (shape[i] == ShapedType::kDynamic)
shape[i] = Torch::kUnknownSize;
}
for (int64_t i = 0; i < rank; ++i) {
if (current[i] == permutations[i])
continue;

12
test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir
View File

@ -968,6 +968,18 @@ func.func @test_transpose_all_permutations_4(%arg0: !torch.vtensor<[2,3,4],f32>)
return %0 : !torch.vtensor<[4,2,3],f32>
}
// -----
// CHECK-LABEL: func.func @test_transpose_dynamic
func.func @test_transpose_dynamic(%arg0: !torch.vtensor<[?,32,5,128],f32>) -> !torch.vtensor<[?,5,32,128],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64} {
// CHECK-DAG: %[[I1:.+]] = torch.constant.int 1
// CHECK-DAG: %[[I2:.+]] = torch.constant.int 2
// CHECK: %[[TRANSPOSE:.+]] = torch.aten.transpose.int %arg0, %[[I1]], %[[I2]] : !torch.vtensor<[?,32,5,128],f32>, !torch.int, !torch.int -> !torch.vtensor<[?,5,32,128],f32>
%0 = torch.operator "onnx.Transpose"(%arg0) {torch.onnx.perm = [0 : si64, 2 : si64, 1 : si64, 3 : si64]} : (!torch.vtensor<[?,32,5,128],f32>) -> !torch.vtensor<[?,5,32,128],f32>
return %0 : !torch.vtensor<[?,5,32,128],f32>
}
// -----
// CHECK-LABEL: func.func @test_slice