[Mlir-commits] [mlir] ed4d12c - [mlir][Shape] Fix a crash when folding nary broadcast ops

[Benjamin Kramer]

Author: Benjamin Kramer
Date: 2021-02-22T20:06:37+01:00
New Revision: ed4d12c2ce30e5a1c270e2b4a08e67851ff8b295

URL: https://github.com/llvm/llvm-project/commit/ed4d12c2ce30e5a1c270e2b4a08e67851ff8b295
DIFF: https://github.com/llvm/llvm-project/commit/ed4d12c2ce30e5a1c270e2b4a08e67851ff8b295.diff

LOG: [mlir][Shape] Fix a crash when folding nary broadcast ops

operands[2] can be nullptr here. I'm not able to build a lit test for
this because of the commutative reordering of operands. It's possible to
trigger this with a createOrFold<BroadcastOp> though.

Differential Revision: https://reviews.llvm.org/D97206

Added: 
    

Modified: 
    mlir/lib/Dialect/Shape/IR/Shape.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 6146239cd4f8..73506525b22c 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -358,7 +358,7 @@ OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
     return nullptr;
 
   // TODO: Support folding with more than 2 input shapes
-  if (operands.size() > 2 && !operands[2].isa<StringAttr>())
+  if (shapes().size() > 2)
     return nullptr;
 
   auto rhsShape = llvm::to_vector<6>(