Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Initial resize scheduler #3556

Merged
merged 38 commits into from
Dec 17, 2024
Merged
Changes from 1 commit
Commits
Show all changes
38 commits
Select commit Hold shift + click to select a range
11f5dce
Always enable IdModel-based indexing when resize is used
naoyam Dec 11, 2024
05ea88f
Don't run the tests without IdModel
naoyam Dec 11, 2024
0ad9fea
fix
naoyam Dec 11, 2024
4f14988
Allocation ordering fix
naoyam Dec 11, 2024
f3ce2d9
Merge remote-tracking branch 'origin/enable_id_model_for_resize' into…
naoyam Dec 11, 2024
0d35147
rotation + residual
naoyam Dec 9, 2024
7934e63
wip
naoyam Dec 10, 2024
9e71bc5
move DomainMap to its own file
naoyam Dec 10, 2024
57600bd
Use the reference finder of pointwise scheduler
naoyam Dec 10, 2024
25ebe94
use the scheduler in the resize test
naoyam Dec 11, 2024
55b8499
WIP
naoyam Dec 11, 2024
791c85b
Merge remote-tracking branch 'origin/main' into resize_scheduler_init…
naoyam Dec 11, 2024
839f23c
cleanup
naoyam Dec 11, 2024
7a10f02
WAR
naoyam Dec 11, 2024
c80dd91
Fix the failed alias test thanks to @wujingyue
naoyam Dec 11, 2024
9167cf0
cleanup
naoyam Dec 11, 2024
df63df2
cleanup
naoyam Dec 11, 2024
52acb42
cleanup
naoyam Dec 11, 2024
6363298
Merge branch 'main' into resize_scheduler_initial_version
naoyam Dec 13, 2024
ca09b93
Merge branch 'main' into resize_scheduler_initial_version
naoyam Dec 13, 2024
0d0a4d6
PR feedback
naoyam Dec 13, 2024
be3aee9
fix
naoyam Dec 13, 2024
4368e80
Rename DomainMap to PointwiseDomainMap
naoyam Dec 13, 2024
2a6f059
Merge remote-tracking branch 'origin/main' into resize_scheduler_init…
naoyam Dec 13, 2024
91e7d3e
Merge remote-tracking branch 'origin/main' into resize_scheduler_init…
naoyam Dec 13, 2024
96ac0fa
merge fix
naoyam Dec 13, 2024
7e9413a
python frontend fix
naoyam Dec 15, 2024
40dd2c2
fix pattern match
naoyam Dec 15, 2024
8056cfa
fix
naoyam Dec 15, 2024
b9415e1
test fix
naoyam Dec 15, 2024
aebfd51
Disable segmentation
naoyam Dec 15, 2024
c264867
fix
naoyam Dec 16, 2024
8e82996
cleanup
naoyam Dec 16, 2024
c4c1136
Disable resize scheduler by default
naoyam Dec 17, 2024
57f2279
format
naoyam Dec 17, 2024
0b8e9ba
format
naoyam Dec 17, 2024
dc4f42d
Merge branch 'main' into resize_scheduler_initial_version
naoyam Dec 17, 2024
0422ce4
Merge branch 'main' into resize_scheduler_initial_version
naoyam Dec 17, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Disable resize scheduler by default
naoyam committed Dec 17, 2024
commit c4c11367554789179aeb58cdddfa45b540ad367b
7 changes: 4 additions & 3 deletions csrc/options.cpp
Original file line number Diff line number Diff line change
@@ -154,15 +154,16 @@ const std::unordered_map<std::string, EnableOption>& getEnableOptions() {
{"fuse_matmul", EnableOption::FuseMatmul},
{"fuse_multiple_matmuls", EnableOption::FuseMultipleMatmuls},
{"id_model", EnableOption::IdModel},
{"io_to_lower_precision", EnableOption::IoToLowerPrecision},
{"kernel_db", EnableOption::KernelDb},
{"kernel_debug", EnableOption::KernelDebug},
{"kernel_lineinfo", EnableOption::KernelLineInfo},
{"kernel_profile", EnableOption::KernelProfile},
{"memory_promotion", EnableOption::MemoryPromotion},
{"reuse_zeroed_memory", EnableOption::ReuseZeroedMemory},
{"resize_scheduler", EnableOption::ResizeScheduler},
{"static_fusion_count", EnableOption::StaticFusionCount},
{"warn_register_spill", EnableOption::WarnRegisterSpill},
{"io_to_lower_precision", EnableOption::IoToLowerPrecision},
{"kernel_debug", EnableOption::KernelDebug},
{"kernel_lineinfo", EnableOption::KernelLineInfo},
};
return available_options;
}
13 changes: 7 additions & 6 deletions csrc/options.h
Original file line number Diff line number Diff line change
@@ -93,17 +93,18 @@ enum class EnableOption {
FuseMatmul, //! Enable automatic fusion of matmul and linear ops
FuseMultipleMatmuls, //! Allow fusing more than one matmul in a single kernel
IdModel, //! Enable IdModel
KernelDb, //! Enable Kernel Database
KernelProfile, //! Enable intra-kernel performance profiling
MemoryPromotion, //! Enable promotion of memory types for non-pointwise ops
StaticFusionCount, //! Enable using single static count in kernel name
ReuseZeroedMemory, //! Re-use zeroed memory used for grid synchronization
WarnRegisterSpill, //! Enable warnings of register spill
IoToLowerPrecision, //! Enable castInputOutputToLowerPrecision. #1889 explains
//! why we disabled it by default.
KernelDb, //! Enable Kernel Database
KernelDebug, //! Enable debug mode in nvrtc
KernelLineInfo, //! Embed line info to compiled kernel, and dump the full CUDA
//! C++ code
KernelProfile, //! Enable intra-kernel performance profiling
MemoryPromotion, //! Enable promotion of memory types for non-pointwise ops
ReuseZeroedMemory, //! Re-use zeroed memory used for grid synchronization
ResizeScheduler, //! Enable the resize scheduler
StaticFusionCount, //! Enable using single static count in kernel name
WarnRegisterSpill, //! Enable warnings of register spill
EndOfOption //! Placeholder for counting the number of elements
};

6 changes: 6 additions & 0 deletions csrc/scheduler/resize.cpp
Original file line number Diff line number Diff line change
@@ -33,6 +33,12 @@ TensorView* getReferenceTensor(Fusion* fusion) {
} // namespace

bool ResizeScheduler::canScheduleCompileTime(Fusion* fusion) {
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In this initial version, I'm trying to make it very restrictive. Will have several follow-up PRs to schedule the whole RoPE module.

if (!isOptionEnabled(EnableOption::ResizeScheduler)) {
scheduler_debug_utils::canScheduleRejectReason(
schedulerType(), "Not enabled");
return false;
}

if (!ir_utils::hasOpsOfType<SliceOp, PadOp>(fusion)) {
scheduler_debug_utils::canScheduleRejectReason(
schedulerType(), "No resize op to schedule");
23 changes: 21 additions & 2 deletions tests/cpp/test_resize.cpp
Copy link
Collaborator Author

@naoyam naoyam Dec 10, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing changed with the tests here (except replacing set with sin and one disabled test) but just extended some of the existing tests to use the resize scheduler as well. Not all patterns are supported yet, so they just call GTEST_SKIP for now.

Original file line number Diff line number Diff line change
@@ -55,8 +55,27 @@ void checkLoopDomainEquivalence(

} // namespace

using ResizeTest = NVFuserTest;
using ResizeSchedulerTest = NVFuserFixtureParamTest<bool>;
class ResizeTest : public NVFuserTest {
protected:
void SetUp() override {
EnableOptionsGuard::getCurOptions().set(EnableOption::ResizeScheduler);
NVFuserTest::SetUp();
}

private:
EnableOptionsGuard enable_options_guard_;
};

class ResizeSchedulerTest : public NVFuserFixtureParamTest<bool> {
protected:
void SetUp() override {
EnableOptionsGuard::getCurOptions().set(EnableOption::ResizeScheduler);
NVFuserFixtureParamTest<bool>::SetUp();
}

private:
EnableOptionsGuard enable_options_guard_;
};

using testing::Each;
using testing::HasSubstr;