vinx13 commented on code in PR #12127:
URL: https://github.com/apache/tvm/pull/12127#discussion_r935907425
##########
src/meta_schedule/schedule_rule/winograd.cc:
##########
@@ -74,6 +80,32 @@ inline LoopRV ScheduleDataPack(Schedule sch, BlockRV block) {
return t1[1];
}
+inline LoopRV ScheduleDataPackNCHW(Schedule sch, BlockRV block) {
+ Array<LoopRV> loops = sch->GetLoops(block);
+ ICHECK_EQ(loops.size(), 6);
+
+ if (const int64_t* i = tir::GetLoopIntExtent(sch->GetSRef(loops[0]))) {
+ if (*i <= 16) {
+ sch->Unroll(loops[0]);
+ }
+ }
+ if (const int64_t* i = tir::GetLoopIntExtent(sch->GetSRef(loops[1]))) {
+ if (*i <= 16) {
+ sch->Unroll(loops[1]);
+ }
+ }
+ sch->Unroll(loops[4]);
+ sch->Unroll(loops[5]);
+
+ Array<ExprRV> factors = sch->SamplePerfectTile(loops[3], /*n=*/2,
/*max_innermost_factor=*/64);
+ Array<LoopRV> split =
+ sch->Split(loops[3], /*factors=*/{factors[0], factors[1]},
/*preserve_unit_loops=*/true);
Review Comment:
```suggestion
sch->Split(loops[3], /*factors=*/factors,
/*preserve_unit_loops=*/true);
```
##########
src/meta_schedule/schedule_rule/winograd.cc:
##########
@@ -92,6 +124,50 @@ TVM_REGISTER_GLOBAL("meta_schedule.winograd_data_pack.llvm")
return {sch};
});
+TVM_REGISTER_GLOBAL("meta_schedule.winograd_output.nchw.cuda")
+ .set_body_typed([](Schedule sch, BlockRV output) -> Array<Schedule> {
+ // get loops
+ Array<LoopRV> loops = sch->GetLoops(output);
+ ICHECK_EQ(loops.size(), 4);
+
+ BlockRV OL{nullptr};
+ if (sch->GetConsumers(output).size() > 0) {
+ OL = GetOnlyConsumer(sch, output);
+ sch->SetScope(OL, /*buffer_index=*/0, /*storage_scope=*/"local");
+ }
+
+ // tile
+ BlockRV inverse = GetOnlyProducer(sch, output);
+ Optional<PrimExpr> tile_size =
+ tir::GetAnn<PrimExpr>(sch->GetSRef(output), "winograd_tile_size");
+ ICHECK(tile_size.defined()) << "Winograd tile size is not defined in
block annotation!";
+ Array<LoopRV> split0 = sch->Split(loops[2], {NullOpt,
tile_size.value()});
+ Array<LoopRV> split1 = sch->Split(loops[3], {NullOpt,
tile_size.value()});
+ sch->Reorder({split0[0], split1[0], split0[1], split1[1]});
+
+ // compute_at
+ sch->ComputeAt(inverse, /*loop_rv=*/split1[0],
+ /*preserve_unit_loops=*/true);
+ if (OL.defined()) {
+ while (sch->GetConsumers(OL).size() > 0) {
+ BlockRV next_OL = GetOnlyConsumer(sch, OL);
Review Comment:
Will consumers be handled automatically by auto_inline rule? If so, this and
`SetScope` are not necessary
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
