From af4aa931aef6e5357441a7d2be07e39575a4c440 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 3 Mar 2025 22:02:06 -0800 Subject: [PATCH 1/3] Make sure IterDomain::resize generates non-symbolic IDs --- csrc/ir/interface_nodes.h | 6 +++++- csrc/ir/internal_base_nodes.h | 6 +++++- csrc/ir/nodes.cpp | 6 ++++-- csrc/scheduler/tools/loop_domain_scheduler.cpp | 8 +++++++- csrc/tensor_view.cpp | 5 +++-- 5 files changed, 24 insertions(+), 7 deletions(-) diff --git a/csrc/ir/interface_nodes.h b/csrc/ir/interface_nodes.h index 56d97018a66..cffd277eeb2 100644 --- a/csrc/ir/interface_nodes.h +++ b/csrc/ir/interface_nodes.h @@ -614,7 +614,11 @@ class NVF_API TensorView : public Val { //! Resize an IterDomain by expanding both the left and right sides //! by given widths. The resulting IterDomain has an extent of //! (left_expansion + axis->extent() + right_expansion). - TensorView* resize(int64_t axis, Val* left_expansion, Val* right_expansion); + TensorView* resize( + int64_t axis, + Val* left_expansion, + Val* right_expansion, + std::optional iter_type = std::nullopt); // WARNING: rFactor does not return this TensorView, ir returns a new // tensorview consumed by this! diff --git a/csrc/ir/internal_base_nodes.h b/csrc/ir/internal_base_nodes.h index bd8e14d9ff3..1094713f201 100644 --- a/csrc/ir/internal_base_nodes.h +++ b/csrc/ir/internal_base_nodes.h @@ -690,7 +690,11 @@ class TensorDomain : public Val { SwizzleMode swizzle_mode = SwizzleMode::Data); // Resize an axis by left_expansion and right_expansion - void resize(int64_t axis, Val* left_expansion, Val* right_expansion); + void resize( + int64_t axis, + Val* left_expansion, + Val* right_expansion, + std::optional iter_type = std::nullopt); // Transform TensorView according to merge and split transformations TensorDomain* view(const AnalyzeViewResult& view_analysis); diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp index 05d89df7e0a..1e3b77973ce 100644 --- a/csrc/ir/nodes.cpp +++ b/csrc/ir/nodes.cpp @@ -3697,13 +3697,15 @@ void TensorDomain::swizzle( void TensorDomain::resize( int64_t axis, Val* left_expansion, - Val* right_expansion) { + Val* right_expansion, + std::optional iter_type) { NVF_ERROR(nDims() > 0, "Tried to do resize on a 0-dim domain"); axis = wrapDim(axis); IterDomain* id = this->axis(axis); - auto resized_id = IterDomain::resize(id, left_expansion, right_expansion); + auto resized_id = + IterDomain::resize(id, left_expansion, right_expansion, false, iter_type); loop_domain_.at(axis) = resized_id; resetDomains(); } diff --git a/csrc/scheduler/tools/loop_domain_scheduler.cpp b/csrc/scheduler/tools/loop_domain_scheduler.cpp index ccab7fce0e1..3f11d7b86cb 100644 --- a/csrc/scheduler/tools/loop_domain_scheduler.cpp +++ b/csrc/scheduler/tools/loop_domain_scheduler.cpp @@ -151,10 +151,16 @@ class ReplayForwardTransformOnLoopDomain : OptInConstDispatch { void handle(const Resize* resize) final { NVF_ERROR(input_loop_ids_.size() == 1); + NVF_ERROR( + resize->out()->getIterType() != IterType::Symbolic, + "Unexpected to have a symbolic ID: ", + resize->out()->toString()); + // Pass the iter type explicitly to avoid generating a symblic ID tv_->resize( getLoopIdPosition(input_loop_ids_.at(0)), resize->leftExpand(), - resize->rightExpand()); + resize->rightExpand(), + resize->out()->getIterType()); } void handle(const Swizzle2D* swizzle_2d) final { diff --git a/csrc/tensor_view.cpp b/csrc/tensor_view.cpp index 8f1c0143f17..d851532dcdd 100644 --- a/csrc/tensor_view.cpp +++ b/csrc/tensor_view.cpp @@ -527,7 +527,8 @@ TensorView* TensorView::merge(int64_t axis_o, int64_t axis_i) { TensorView* TensorView::resize( int64_t axis, Val* left_expansion, - Val* right_expansion) { + Val* right_expansion, + std::optional iter_type) { NVF_ERROR( nDims() > 0, "Tried to do resize on a 0-dim TensorView. ", @@ -560,7 +561,7 @@ TensorView* TensorView::resize( " Parallelization strategy must be set after calling resize: ", toString()); - domain()->resize(axis, left_expansion, right_expansion); + domain()->resize(axis, left_expansion, right_expansion, iter_type); return this; } From 6300b1e774f494a2e2f56346690fa60951229628 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Wed, 5 Mar 2025 12:45:54 -0800 Subject: [PATCH 2/3] comment --- csrc/ir/nodes.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp index 1e3b77973ce..574a10452c0 100644 --- a/csrc/ir/nodes.cpp +++ b/csrc/ir/nodes.cpp @@ -3704,8 +3704,12 @@ void TensorDomain::resize( IterDomain* id = this->axis(axis); - auto resized_id = - IterDomain::resize(id, left_expansion, right_expansion, false, iter_type); + auto resized_id = IterDomain::resize( + id, + left_expansion, + right_expansion, + /*mark_as_rfactor=*/false, + iter_type); loop_domain_.at(axis) = resized_id; resetDomains(); } From 809b5836b1e7cc3e0992b33e18f3a1b977c3b457 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Thu, 6 Mar 2025 08:40:04 -0800 Subject: [PATCH 3/3] Just to trigger CI