diff --git a/csrc/id_model/indexing.cpp b/csrc/id_model/indexing.cpp index 4c42da75010..9d71a0c3e3e 100644 --- a/csrc/id_model/indexing.cpp +++ b/csrc/id_model/indexing.cpp @@ -34,6 +34,8 @@ namespace nvfuser { TensorIndexer::TensorIndexer(IdModel& id_model) : id_model_(id_model) { + NVF_ERROR(isSupported(id_model.fusion())); + buildLoopIndexMap(); if (isDebugDumpEnabled(DebugDumpOption::IndexingVerbose)) { diff --git a/csrc/scheduler/expr_eval_sched.cpp b/csrc/scheduler/expr_eval_sched.cpp index f2b6ebce41b..f540e545bb5 100644 --- a/csrc/scheduler/expr_eval_sched.cpp +++ b/csrc/scheduler/expr_eval_sched.cpp @@ -64,6 +64,7 @@ bool ExprEvalScheduler::canScheduleCompileTime(Fusion* fusion) { // TODO: remove IndexPutAccumulateOp if (exprs.front() ->isOneOf< + GatherOp, ScatterOp, SdpaFwdOp, SdpaBwdOp, diff --git a/csrc/scheduler/registry.cpp b/csrc/scheduler/registry.cpp index 7e47cc253b7..a04714e22b0 100644 --- a/csrc/scheduler/registry.cpp +++ b/csrc/scheduler/registry.cpp @@ -64,6 +64,16 @@ bool checkCanSchedule(Fusion* fusion, SchedulerType scheduler_type) { return false; } + // Support of non-exact gather was dropped when the legacy indexer was + // deprecated + if (std::ranges::any_of( + ir_utils::getOpsOfType(fusion), + [](GatherOp* gather) { return !gather->exactSizes(); })) { + scheduler_debug_utils::canScheduleRejectReason( + scheduler_type, "Non-exact gather ops"); + return false; + } + // Fusions with `MatmulOp, LinearOp, MmaOp` can only be accepted by Matmul // scheduler. if (scheduler_type != SchedulerType::Matmul && diff --git a/tests/cpp/test_gather.cpp b/tests/cpp/test_gather.cpp index 0924b33a70b..92e78f7c4fa 100644 --- a/tests/cpp/test_gather.cpp +++ b/tests/cpp/test_gather.cpp @@ -583,7 +583,7 @@ TEST_F(GatherTest, TakeAlongAxisIntermediateTensorReduction1) { validateSegmentation( executor_cache.getMostRecentKernelRuntime(), - {SchedulerType::Reduction, SchedulerType::PointWise}); + {SchedulerType::Reduction, SchedulerType::ExprEval}); testValidate(&fusion, outputs, {t0, t1}, __LINE__, __FILE__); } @@ -1128,7 +1128,8 @@ TEST_F(GatherTest, TakeAlongAxisCrossEntropyLoss) { } // Test grouped reduction on IterType::GatherScatter -TEST_F(GatherTest, GatherIterGoupedReduction) { +// Codegen support of non-exact gather dropped +TEST_F(GatherTest, DISABLED_GatherIterGoupedReduction) { const int max_dim_size = 128; auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0); auto options_i = at::TensorOptions().dtype(at::kLong).device(at::kCUDA, 0); @@ -1212,7 +1213,8 @@ TEST_F(GatherTest, GatherIterGoupedReduction) { lparams); } -TEST_F(GatherTest, SameTvUsedAsLookupAndIndex) { +// Codegen support of non-exact gather dropped +TEST_F(GatherTest, DISABLED_SameTvUsedAsLookupAndIndex) { auto fusion_ptr = std::make_unique(); Fusion& fusion = *fusion_ptr.get(); FusionGuard fg(&fusion);