Patches

lib/Dialect/TTIR/Transforms/Fusion.cpp

@@ -128,13 +128,26 @@ class FuseSoftmax : public OpRewritePattern<ttir::DivOp> {
 
     // There may be a broadcast op between the sum op and the div op.
     if (!sumOp && dyn_cast_or_null<ttir::BroadcastOp>(rhs.getDefiningOp())) {
+      auto broadcastOp =
+          dyn_cast_or_null<ttir::BroadcastOp>(rhs.getDefiningOp());
+
+      // Broadcast op should only have one use.
+      if (!broadcastOp->hasOneUse()) {
+        return failure();
+      }
+
       sumOp = dyn_cast_or_null<ttir::SumOp>(
-          rhs.getDefiningOp()->getOperand(0).getDefiningOp());
+          broadcastOp->getOperand(0).getDefiningOp());
     }
     if (!expOp || !sumOp) {
       return failure();
     }
 
+    // Both the exp and sum op should have only one use.
+    if (!expOp->hasOneUse() != 1 || !sumOp->hasOneUse() != 1) {
+      return failure();
+    }
+
     if (sumOp.getInput() != expOp.getResult(0)) {
       return failure();
     }