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

[Torch Dialect] fix torch.uint8's dtype infer (#2227)

pull/2231/head snapshot-20230613.868
Yuanqiang Liu 2023-06-13 10:38:20 +08:00 committed by GitHub
parent dd5992514d
commit ddea56a832
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
lib/Dialect/Torch/Utils/Utils.cpp
View File

@ -103,6 +103,7 @@ Torch::getTypeForScalarType(MLIRContext *context,
case torch_upstream::ScalarType::Half:
return mlir::FloatType::getF16(context);
case torch_upstream::ScalarType::Byte:
return mlir::IntegerType::get(context, 8, mlir::IntegerType::Unsigned);
case torch_upstream::ScalarType::Char:
return mlir::IntegerType::get(context, 8, signedness);
case torch_upstream::ScalarType::ComplexHalf:

10
test/Dialect/Torch/torch-function-to-torch-backend-pipeline.mlir
View File

@ -15,3 +15,13 @@ func.func @torch.aten.argmax(%arg0: !torch.vtensor<[?,?],f32>) -> !torch.vtensor
%0 = torch.aten.argmax %arg0, %int0, %true : !torch.vtensor<[?,?],f32>, !torch.int, !torch.bool -> !torch.vtensor<[1,?],si64>
return %0 : !torch.vtensor<[1,?],si64>
}
// CHECK-LABEL: func.func @torch.uint8
func.func @torch.uint8(%arg0: !torch.tensor {torch.type_bound = !torch.vtensor<[3,4],ui8>}) -> !torch.tensor {
%int12 = torch.constant.int 12
%0 = torch.prim.ListConstruct %int12 : (!torch.int) -> !torch.list<int>
// CHECK: torch.aten.view
// CHECK-SAME: !torch.vtensor<[12],ui8>
%1 = torch.aten.reshape %arg0, %0 : !torch.tensor, !torch.list<int> -> !torch.tensor
return %1 : !torch.tensor
}