[CINN] Fix TileBroadcastTactic matching conditions (#72348)

Enigmatisms · web-flow · commit bb59a58465e9 · 2025-04-18T14:18:57.000+08:00
diff --git a/paddle/cinn/ir/group_schedule/tactic/tile_broadcast_tactic.cc b/paddle/cinn/ir/group_schedule/tactic/tile_broadcast_tactic.cc
@@ -206,6 +206,7 @@ bool IsElementwiseOrBroadcast(const ir::Store& dst, const ir::Load& src) {
 
 bool CheckAllElementwiseOrBroadcast(ir::IRSchedule* sch) {
   for (auto& block : sch->GetAllBlocks()) {
+    if (ir::analyzer::IsReductionSBlock(block)) return false;
     ir::Expr store = ir::analyzer::GetStoreOfSBlock(block);
     auto* store_node = store.As<ir::Store>();
     for (auto& load : CollectLoads(store_node->value)) {
@@ -564,6 +565,7 @@ void TileBroadcastTactic::Apply(ir::IRSchedule* sch,
     block_size = CalcNumWarps(preserved_size_ >> 5);
     if (block_size == -1) {
       applied_layout_ = BroadcastLayout::Invalid;
+      return;
     }
     block_size = std::clamp(block_size << 5, 128, 1024);
   }
diff --git a/test/dygraph_to_static/test_pir_selectedrows.py b/test/dygraph_to_static/test_pir_selectedrows.py
@@ -17,7 +17,6 @@
 
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    test_phi_only,
 )
 
 import paddle
@@ -74,7 +73,6 @@ def forward_static():
 
 
 class TestSimnet(Dy2StTestBase):
-    @test_phi_only
     def test_dygraph_static_same_loss(self):
         dygraph_value = forward_dygraph()
         static_value = forward_static()

Original file line number	Diff line number	Diff line change
`@@ -206,6 +206,7 @@ bool IsElementwiseOrBroadcast(const ir::Store& dst, const ir::Load& src) {`
`206`	`206`
`207`	`207`	`bool CheckAllElementwiseOrBroadcast(ir::IRSchedule* sch) {`
`208`	`208`	`for (auto& block : sch->GetAllBlocks()) {`
	`209`	`+ if (ir::analyzer::IsReductionSBlock(block)) return false;`
`209`	`210`	`ir::Expr store = ir::analyzer::GetStoreOfSBlock(block);`
`210`	`211`	`auto* store_node = store.As<ir::Store>();`
`211`	`212`	`for (auto& load : CollectLoads(store_node->value)) {`
`@@ -564,6 +565,7 @@ void TileBroadcastTactic::Apply(ir::IRSchedule* sch,`
`564`	`565`	`block_size = CalcNumWarps(preserved_size_ >> 5);`
`565`	`566`	`if (block_size == -1) {`
`566`	`567`	`applied_layout_ = BroadcastLayout::Invalid;`
	`568`	`+ return;`
`567`	`569`	`}`
`568`	`570`	`block_size = std::clamp(block_size << 5, 128, 1024);`
`569`	`571`	`}`