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

[stablehlo] add dtype conversion when converting AtenScalarImplicitOp (#2439)

pull/2441/head
Jiawei Wu 2023-09-06 01:57:15 +08:00 committed by GitHub
parent 3841fe3035
commit c93c6970e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
lib/Conversion/TorchToStablehlo/Basic.cpp
View File

@ -898,7 +898,16 @@ template <>
LogicalResult ConvertAtenOp<AtenScalarImplicitOp>::matchAndRewrite( LogicalResult ConvertAtenOp<AtenScalarImplicitOp>::matchAndRewrite(
AtenScalarImplicitOp op, OpAdaptor adaptor, AtenScalarImplicitOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const { ConversionPatternRewriter &rewriter) const {
rewriter.replaceOpWithNewOp<tensor::ExtractOp>(op, adaptor.getA()); Location loc = op.getLoc();
Type inputDtype =
op.getA().getType().template cast<BaseTensorType>().getDtype();
Type resultType =
this->getTypeConverter()->convertType(op->getResult(0).getType());
auto result =
rewriter.create<tensor::ExtractOp>(loc, adaptor.getA());
rewriter.replaceOp(
op, convertScalarToDtype(rewriter, loc, result, resultType, inputDtype));
return success(); return success();
} }