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

Allow lower bounds check in broadcast simplification

numpy_style_broadcast
Quinn Dawkins 2023-08-27 21:53:37 -04:00
parent 83dd551656
commit dce6c83f96
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
externals/llvm-external-projects/torch-mlir-dialects/lib/Dialect/TMTensor/Transforms/ConvertBroadcastToLinalg.cpp vendored
View File

@ -53,6 +53,14 @@ public:
broadcastedStatus.push_back(false);
continue;
}
// If the dim is non-unit, then it is an assert + non-broadcasted dim.
FailureOr<int64_t> lb = ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::LB, input, i);
if (succeeded(lb) && *lb > 1) {
// TODO: Insert appropriate assert here.
broadcastedStatus.push_back(false);
continue;
}
FailureOr<bool> isUnit =
ValueBoundsConstraintSet::areEqual(input, oneIndex, i, std::nullopt);
if (succeeded(isUnit) && *isUnit) {