@@ -809,20 +809,21 @@ static bool isAttentionMaskGenerator(Operation *op) {
809
809
return false ;
810
810
}
811
811
812
- static bool hasFusableOrLowMemoryUsers (Operation *op) {
813
- bool hasMemoryHungryNonFusableConsumer = false ;
812
+ static bool hasExplicitNonFusableUsers (Operation *op) {
813
+ bool hasNonFusableUse = false ;
814
814
for (Operation *user : op->getUsers ()) {
815
815
if (isa<IREE::LinalgExt::LinalgFusionOpInterface>(user))
816
816
continue ;
817
- // TODO: Calling iree_linalg_ext.scan ops "memory-hungry" is a bit of a
818
- // stretch - the actual issue comes down to poor support in tiling configs,
819
- // leading to overly large stack-bound allocations. Backend compilers should
820
- // often cope with fusing the resulting loops.
821
- // So improve the tiling logic for ScanOp's, while also considering ways to
822
- // express simple fusions within this cumulative reduction intrinsic.
823
- hasMemoryHungryNonFusableConsumer |= isa<IREE::LinalgExt::ScanOp>(user);
824
- }
825
- return !hasMemoryHungryNonFusableConsumer;
817
+ // TODO: The issue with iree_linalg_ext.scan compared to other non-fusable
818
+ // ops comes down to poor support in tiling configs, leading to overly large
819
+ // stack-bound allocations. In practice, backend compilers should often cope
820
+ // with fusing the resulting loops even without Linalg-level fusion.
821
+ // So long-term, we improve the tiling logic for ScanOp's, while also
822
+ // considering ways to express simple fusions within this cumulative
823
+ // reduction intrinsic.
824
+ hasNonFusableUse |= isa<IREE::LinalgExt::ScanOp>(user);
825
+ }
826
+ return hasNonFusableUse;
826
827
}
827
828
828
829
// / Operations that are cloned into dispatch regions formed with other
@@ -842,8 +843,8 @@ bool isClonableIntoDispatchOp(Operation *op,
842
843
return true ;
843
844
}
844
845
// TODO: Tune the cases excluded through hasFusableOrLowMemoryUsers
845
- // condition in a more targetted manner, then remove the condition.
846
- if (isa<linalg::LinalgOp>(op) && ! hasFusableOrLowMemoryUsers (op)) {
846
+ // condition in a more targeted manner, then remove the condition.
847
+ if (isa<linalg::LinalgOp>(op) && hasExplicitNonFusableUsers (op)) {
847
848
return false ;
848
849
}
849
850
if (LinalgExt::isBitExtendOp (op)) {
0 commit comments