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

[stablehlo] verify stablehlo backend contract (#3338)

pull/3373/head
penguin_wwy 2024-05-16 11:03:43 +08:00 committed by GitHub
parent 0ca88028cd
commit 405f884522
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
lib/Dialect/TorchConversion/Transforms/VerifyStablehloBackendContract.cpp
View File

@ -11,10 +11,9 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Transforms/DialectConversion.h"
#include "stablehlo/dialect/ChloOps.h"
#include "stablehlo/dialect/StablehloOps.h"
@ -47,13 +46,21 @@ class VerifyStablehloBackendContractPass
// Structural operations.
target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(
opHasLegalTypes);
// Shape operations.
target.addDynamicallyLegalOp<shape::ShapeOfOp>(opHasLegalTypes);
target.addLegalDialect<chlo::ChloDialect>();
target.addLegalDialect<stablehlo::StablehloDialect>();
target.addLegalDialect<tensor::TensorDialect>();
target.addLegalDialect<arith::ArithDialect>();
target.addLegalDialect<math::MathDialect>();
target.addLegalDialect<shape::ShapeDialect>();
auto moduleOp = getOperation();
RewritePatternSet patterns(context);
if (failed(applyFullConversion(moduleOp, target, std::move(patterns)))) {
emitError(moduleOp.getLoc())
<< "Module does not conform to the Stablehlo backend contract. "
"See dialect conversion legality information above.";
return signalPassFailure();
}
}
};
} // namespace