colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[MLIR][MHLO] Implement shape reification for `dynamic_broadcast_in_dim` 

PiperOrigin-RevId: 363622714
This commit is contained in:
A. Unique TensorFlower 2021-03-18 03:38:18 -07:00 committed by TensorFlow MLIR Team
parent aad6d39952
commit 0c4a89e52c
2 changed files with 15 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
View File

@ -825,8 +825,9 @@ def HLO_BroadcastInDimOp : HLO_Op<"broadcast_in_dim",
let hasCustomHLOConverter = 1;
}
def HLO_DynamicBroadcastInDimOp : HLO_Op<"dynamic_broadcast_in_dim",
[NoSideEffect]> {
def HLO_DynamicBroadcastInDimOp : HLO_Op<"dynamic_broadcast_in_dim", [
NoSideEffect, DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
["reifyReturnTypeShapes"]>]> {
string summary = "Broadcast a tensor into the given dynamic shape by adding dimensions.";
string description = [{
This is a generalization of the BroadcastInDimOp which accepts its output

12
lib/Dialect/mhlo/IR/hlo_ops.cc
View File

@ -908,6 +908,18 @@ void DynamicBroadcastInDimOp::getCanonicalizationPatterns(
context);
}
LogicalResult DynamicBroadcastInDimOp::inferReturnTypeComponents(
MLIRContext*, llvm::Optional<mlir::Location>, ValueRange, DictionaryAttr,
RegionRange, llvm::SmallVectorImpl<mlir::ShapedTypeComponents>&) {
return failure();
}
LogicalResult DynamicBroadcastInDimOp::reifyReturnTypeShapes(
OpBuilder&, SmallVectorImpl<Value>& reifiedReturnShapes) {
reifiedReturnShapes.push_back(output_dimensions());
return success();
}
//===----------------------------------------------------------------------===//
// ClampOp
//===----------------------------------------------------------------------===//