From 98484b94179f37c9448019909dd1a3f873a02b78 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Wed, 22 Feb 2023 15:20:51 -0800 Subject: [PATCH 1/3] [TIR] Allow TransformLayout with non-inversible index map TransformLayout requires the index map to have inverse map that can be calculated by the analyzer in order to check whether padding is added. However, such check doesn't always work for all cases because of limitation of the affine analysis that can only handle a set of supported patterns. In some cases, even if the index map doesn't introduce padding, the schedule primitive throws `TransformationIntroducesPaddingError` because it fails to calculate the inverse index map. It is safe to allow buffer being padded without providing pad_value because the original loop extent is not changed and the padded region is not accessed. This PR changes the behavior of `TransformLayout` to allow non-inversible index map. Previous discussion: https://discuss.tvm.apache.org/t/conflict-free-shared-memory-permutation-in-tensorir/13959/9 --- .../primitive/layout_transformation.cc | 75 +++++-------------- .../test_tir_schedule_transform_layout.py | 13 +++- 2 files changed, 28 insertions(+), 60 deletions(-) diff --git a/src/tir/schedule/primitive/layout_transformation.cc b/src/tir/schedule/primitive/layout_transformation.cc index 0e993d06dcf1..6af69ab64a6c 100644 --- a/src/tir/schedule/primitive/layout_transformation.cc +++ b/src/tir/schedule/primitive/layout_transformation.cc @@ -92,11 +92,17 @@ class TransformLayoutPlanner : private StmtExprVisitor { std::variant; static TransformPlan Plan(Block block, Buffer old_buffer, Buffer new_buffer, IndexMap index_map, - IndexMap inverse, PrimExpr padding_predicate, Optional pad_value) { ICHECK(!pad_value.defined() || pad_value.value()->final_indices.size() == 1) << "Internal error: Should be caught by ScheduleError checks prior to this point"; TransformLayoutPlanner visitor(old_buffer); + auto [inverse, padding_predicate] = [&]() { + Array region; + for (const auto& dim : old_buffer->shape) { + region.push_back(Range::FromMinExtent(make_zero(dim.dtype()), dim)); + } + return index_map.NonSurjectiveInverse(region); + }(); visitor(block); return visitor.Finalize(new_buffer, index_map, inverse, padding_predicate, pad_value); } @@ -748,15 +754,18 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { * \param old_buffer The target buffer before transformation * \param new_buffer The new buffer after transformation * \param index_map The transformation applied to the buffer + * \param pad_value The value to be used for padding * \return The new AST rooting at the original parent scope and the map from the old block to the * new block */ - static std::pair> Rewrite( - const Block& scope_stmt, const Buffer& old_buffer, const Buffer& new_buffer, - const IndexMap& index_map, const IndexMap& inverse, const PrimExpr& padding_predicate, - const Optional& pad_value) { - auto plan = TransformLayoutPlanner::Plan(scope_stmt, old_buffer, new_buffer, index_map, inverse, - padding_predicate, pad_value); + static std::pair> Rewrite(const Block& scope_stmt, + const Buffer& old_buffer, + const Buffer& new_buffer, + const IndexMap& index_map, + const Optional& pad_value) { + auto plan = pad_value.defined() ? TransformLayoutPlanner::Plan(scope_stmt, old_buffer, + new_buffer, index_map, pad_value) + : TransformLayoutPlanner::NoPaddingRequired{}; arith::Analyzer analyzer; TransformLayoutRewriter rewriter(old_buffer, new_buffer, index_map, plan, &analyzer); @@ -1049,40 +1058,6 @@ class TransformationPaddingExpressionError : public ScheduleError { BufferLoad illegal_load_; }; -class TransformationIntroducesPaddingError : public ScheduleError { - public: - TransformationIntroducesPaddingError(IRModule mod, Buffer buffer, IndexMap index_map, - PrimExpr padding_predicate) - : mod_(std::move(mod)), - buffer_(std::move(buffer)), - index_map_(std::move(index_map)), - padding_predicate_(std::move(padding_predicate)) {} - - String FastErrorString() const final { - std::ostringstream ss; - ss << "ScheduleError: Transformation would introduce padding at " << padding_predicate_ << "."; - return ss.str(); - } - - String DetailRenderTemplate() const final { - auto new_shape = index_map_->MapShape(buffer_->shape); - std::ostringstream os; - os << "The transformation " << index_map_ << " applied on buffer " << buffer_->name - << " of shape " << buffer_->shape << " would result in shape " << new_shape - << ". However, this would introduce padding wherever " << padding_predicate_ << " is true."; - return os.str(); - } - - IRModule mod() const final { return mod_; } - Array LocationsOfInterest() const final { return {}; } - - private: - IRModule mod_; - Buffer buffer_; - IndexMap index_map_; - PrimExpr padding_predicate_; -}; - // Make the dtypes of indices in IndexMap be the same as the dtype of the buffer shape, to avoid // dtype-mismatch issues later. IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const Array& args) { @@ -1147,28 +1122,14 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ : GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); const BlockNode* scope_block = TVM_SREF_TO_BLOCK(scope_sref); - auto [inverse, padding_predicate] = [&]() { - Array region; - for (const auto& dim : old_buffer->shape) { - region.push_back(Range::FromMinExtent(make_zero(dim.dtype()), dim)); - } - return index_map.NonSurjectiveInverse(region); - }(); - - bool has_padding = !is_zero(padding_predicate); - if (has_padding && !pad_value.defined()) { - throw TransformationIntroducesPaddingError(self->mod, old_buffer, index_map, padding_predicate); - } - // Step 2: Infer the shape of the new buffer Buffer new_buffer = old_buffer; new_buffer.CopyOnWrite()->shape = index_map->MapShape(old_buffer->shape); // Step 3: Rewrite BufferLoad/BufferStore access indices, block read/write regions, and block // alloc_buffers. - auto [new_stmt, block_sref_reuse] = - TransformLayoutRewriter::Rewrite(GetRef(scope_block), old_buffer, new_buffer, - index_map, inverse, padding_predicate, pad_value); + auto [new_stmt, block_sref_reuse] = TransformLayoutRewriter::Rewrite( + GetRef(scope_block), old_buffer, new_buffer, index_map, pad_value); Block new_scope_block = Downcast(new_stmt); // Step 4: Rewrite buffer_map of the PrimFunc if necessary. diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index d866de33f100..1561d962b3c8 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -565,8 +565,10 @@ def expected(): A[i // 4, i % 4] = 0 -class TestErrorIfPaddingForbidden(BasePaddingCompare): - """Unless padding is explicitly enabled, should raise error""" +class TestImplicitPadding(BasePaddingCompare): + """When pad_value is None, the buffer can be implicitly padded. The padded region is not + accessed because the original loop extent is not changed. + """ def before(): A = T.alloc_buffer(14, "int32") @@ -575,7 +577,12 @@ def before(): vi = T.axis.remap("S", [i]) A[vi] = 0 - expected = tvm.tir.schedule.schedule.ScheduleError + def expected(): + A = T.alloc_buffer([4, 4], "int32") + for i in T.serial(14): + with T.block("block"): + vi = T.axis.remap("S", [i]) + A[vi // 4, vi % 4] = 0 class TestErrorOnWrongPaddingType(BasePaddingCompare): From 1a68d7a81b25171bd61860a65260d5c36b3e4d04 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 27 Feb 2023 14:01:14 -0800 Subject: [PATCH 2/3] add assume_injective_transform option --- include/tvm/tir/schedule/schedule.h | 9 +- python/tvm/tir/schedule/schedule.py | 18 +++- src/tir/schedule/concrete_schedule.cc | 5 +- src/tir/schedule/concrete_schedule.h | 3 +- src/tir/schedule/primitive.h | 7 +- .../primitive/layout_transformation.cc | 95 ++++++++++++++----- src/tir/schedule/schedule.cc | 4 +- src/tir/schedule/traced_schedule.cc | 9 +- src/tir/schedule/traced_schedule.h | 3 +- .../test_tir_schedule_transform_layout.py | 33 ++++++- 10 files changed, 145 insertions(+), 41 deletions(-) diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index 288601d1cccc..cdc22a294d00 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -642,10 +642,17 @@ class ScheduleNode : public runtime::Object { * Algebraic symplifications, branch elimination, and other * optimizations may assume that this precondition is met, and * may result in incorrect results being returned. + * + * \param assume_injective_transform If set to true, the schedule primitive will assume the + * index_map is injective and skip checking overlapping of the mapped indices. This can be useful + * for complicated index_map that the analysis does not cover. It is the callers' responsibility + * to ensure the index map is injective, otherwise, the correctness of the schedule is not + * guaranteed. */ virtual void TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map, - const Optional& pad_value = NullOpt) = 0; + const Optional& pad_value = NullOpt, + bool assume_injective_transform = false) = 0; /*! * \brief Apply a transformation represented by IndexMap to block diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 4727b4a5f00e..5458f6e8bab6 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -2575,7 +2575,6 @@ def _normalize_buffer_arg( buffer: Union[Tuple[str, int], int, str, Buffer], required_buffer_type=None, ) -> Tuple[str, int, Buffer]: - block_obj: Block = self.get(block) block_name = block_obj.name_hint @@ -2645,6 +2644,8 @@ def transform_layout( buffer: Union[Tuple[str, int], str, Buffer], index_map: Union[IndexMap, Callable], pad_value: Optional[Union[int, float, PrimExpr, IndexMap, Callable]] = None, + *, + assume_injective_transform=False, ) -> None: """Apply a transformation represented by IndexMap to buffer @@ -2711,6 +2712,13 @@ def transform_layout( value to be present in the padding in terms of the transformed index. + assume_injective_transform : bool + + If set to true, the schedule primitive will assume the index_map is injective and skip + checking overlapping of the mapped indices. This can be useful for complicated index_map + that the analysis does not cover. It is the callers' responsibility to ensure the + index map is injective, otherwise, the correctness of the schedule is not guaranteed. + Examples -------- Before transform_layout, in TensorIR, the IR is: @@ -2787,7 +2795,13 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> buffer_index_type_enum = 0 if buffer_index_type == "read" else 1 _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member - self, block, buffer_index, buffer_index_type_enum, index_map, pad_value + self, + block, + buffer_index, + buffer_index_type_enum, + index_map, + pad_value, + assume_injective_transform, ) if axis_separators: _ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index b6af22263e41..8af39b24fdb8 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -800,14 +800,15 @@ void ConcreteScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann void ConcreteScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map, - const Optional& pad_value) { + const Optional& pad_value, + bool assume_injective_transform) { TVM_TIR_SCHEDULE_BEGIN(); auto f_subst = [&](const Var& var) -> Optional { return Downcast>(symbol_table_.Get(var)); }; auto new_index_map = Substitute(index_map, f_subst); tir::TransformLayout(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type, - new_index_map, pad_value); + new_index_map, pad_value, assume_injective_transform); this->state_->DebugVerify(); TVM_TIR_SCHEDULE_END("transform_layout", this->error_render_level_); } diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index 44d9e9b69c94..41168fb016f3 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -148,7 +148,8 @@ class ConcreteScheduleNode : public ScheduleNode { void Unannotate(const BlockRV& block_rv, const String& ann_key) override; /******** Schedule: Layout transformation ********/ void TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, - const IndexMap& index_map, const Optional& pad_value) override; + const IndexMap& index_map, const Optional& pad_value, + bool assume_injective_transform = false) override; void TransformBlockLayout(const BlockRV& block_rv, const IndexMap& index_map) override; void SetAxisSeparator(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index dbc4e235965c..65e6834cac81 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -501,10 +501,15 @@ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const String& * \param buffer_index_type The type of the buffer index, kRead or kWrite. * \param index_map The transformation to apply. * \param pad_value The value to write into padding introduced by the transformation. + * \param assume_injective_transform If set to true, the schedule primitive will assume the + * index_map is injective and skip checking overlapping of the mapped indices. This can be useful + * for complicated index_map that the analysis does not cover. It is the callers' responsibility + * to ensure the index map is injective, otherwise, the correctness of the schedule is not + * guaranteed. */ TVM_DLL void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map, - const Optional& pad_value); + const Optional& pad_value, bool assume_injective_transform); /*! * \brief Apply a transformation represented by IndexMap to block diff --git a/src/tir/schedule/primitive/layout_transformation.cc b/src/tir/schedule/primitive/layout_transformation.cc index 6af69ab64a6c..7eaca74100eb 100644 --- a/src/tir/schedule/primitive/layout_transformation.cc +++ b/src/tir/schedule/primitive/layout_transformation.cc @@ -92,17 +92,11 @@ class TransformLayoutPlanner : private StmtExprVisitor { std::variant; static TransformPlan Plan(Block block, Buffer old_buffer, Buffer new_buffer, IndexMap index_map, + IndexMap inverse, PrimExpr padding_predicate, Optional pad_value) { ICHECK(!pad_value.defined() || pad_value.value()->final_indices.size() == 1) << "Internal error: Should be caught by ScheduleError checks prior to this point"; TransformLayoutPlanner visitor(old_buffer); - auto [inverse, padding_predicate] = [&]() { - Array region; - for (const auto& dim : old_buffer->shape) { - region.push_back(Range::FromMinExtent(make_zero(dim.dtype()), dim)); - } - return index_map.NonSurjectiveInverse(region); - }(); visitor(block); return visitor.Finalize(new_buffer, index_map, inverse, padding_predicate, pad_value); } @@ -754,18 +748,17 @@ class TransformLayoutRewriter : private arith::IRMutatorWithAnalyzer { * \param old_buffer The target buffer before transformation * \param new_buffer The new buffer after transformation * \param index_map The transformation applied to the buffer - * \param pad_value The value to be used for padding * \return The new AST rooting at the original parent scope and the map from the old block to the * new block */ - static std::pair> Rewrite(const Block& scope_stmt, - const Buffer& old_buffer, - const Buffer& new_buffer, - const IndexMap& index_map, - const Optional& pad_value) { - auto plan = pad_value.defined() ? TransformLayoutPlanner::Plan(scope_stmt, old_buffer, - new_buffer, index_map, pad_value) - : TransformLayoutPlanner::NoPaddingRequired{}; + static std::pair> Rewrite( + const Block& scope_stmt, const Buffer& old_buffer, const Buffer& new_buffer, + const IndexMap& index_map, const Optional& opt_inverse, + const PrimExpr& padding_predicate, const Optional& pad_value) { + auto plan = pad_value.defined() ? TransformLayoutPlanner::Plan( + scope_stmt, old_buffer, new_buffer, index_map, + opt_inverse.value(), padding_predicate, pad_value) + : TransformLayoutPlanner::NoPaddingRequired(); arith::Analyzer analyzer; TransformLayoutRewriter rewriter(old_buffer, new_buffer, index_map, plan, &analyzer); @@ -1058,6 +1051,40 @@ class TransformationPaddingExpressionError : public ScheduleError { BufferLoad illegal_load_; }; +class TransformationIntroducesPaddingError : public ScheduleError { + public: + TransformationIntroducesPaddingError(IRModule mod, Buffer buffer, IndexMap index_map, + PrimExpr padding_predicate) + : mod_(std::move(mod)), + buffer_(std::move(buffer)), + index_map_(std::move(index_map)), + padding_predicate_(std::move(padding_predicate)) {} + + String FastErrorString() const final { + std::ostringstream ss; + ss << "ScheduleError: Transformation would introduce padding at " << padding_predicate_ << "."; + return ss.str(); + } + + String DetailRenderTemplate() const final { + auto new_shape = index_map_->MapShape(buffer_->shape); + std::ostringstream os; + os << "The transformation " << index_map_ << " applied on buffer " << buffer_->name + << " of shape " << buffer_->shape << " would result in shape " << new_shape + << ". However, this would introduce padding wherever " << padding_predicate_ << " is true."; + return os.str(); + } + + IRModule mod() const final { return mod_; } + Array LocationsOfInterest() const final { return {}; } + + private: + IRModule mod_; + Buffer buffer_; + IndexMap index_map_; + PrimExpr padding_predicate_; +}; + // Make the dtypes of indices in IndexMap be the same as the dtype of the buffer shape, to avoid // dtype-mismatch issues later. IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const Array& args) { @@ -1094,7 +1121,7 @@ IndexMap LegalizeIndexMapDType(const IndexMap& index_map, const Array& void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map_orig, - const Optional& pad_value) { + const Optional& pad_value, bool assume_injective_transform) { // Step 1: Input handling and error checking const BlockNode* block_ptr = TVM_SREF_TO_BLOCK(block_sref); Buffer old_buffer = @@ -1122,14 +1149,32 @@ void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_ : GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); const BlockNode* scope_block = TVM_SREF_TO_BLOCK(scope_sref); + Optional opt_inverse = NullOpt; + PrimExpr padding_predicate = Bool(false); + if (!assume_injective_transform) { + std::tie(opt_inverse, padding_predicate) = [&]() { + Array region; + for (const auto& dim : old_buffer->shape) { + region.push_back(Range::FromMinExtent(make_zero(dim.dtype()), dim)); + } + return index_map.NonSurjectiveInverse(region); + }(); + } + + bool has_padding = !is_zero(padding_predicate); + if (has_padding && !pad_value.defined()) { + throw TransformationIntroducesPaddingError(self->mod, old_buffer, index_map, padding_predicate); + } + // Step 2: Infer the shape of the new buffer Buffer new_buffer = old_buffer; new_buffer.CopyOnWrite()->shape = index_map->MapShape(old_buffer->shape); // Step 3: Rewrite BufferLoad/BufferStore access indices, block read/write regions, and block // alloc_buffers. - auto [new_stmt, block_sref_reuse] = TransformLayoutRewriter::Rewrite( - GetRef(scope_block), old_buffer, new_buffer, index_map, pad_value); + auto [new_stmt, block_sref_reuse] = + TransformLayoutRewriter::Rewrite(GetRef(scope_block), old_buffer, new_buffer, + index_map, opt_inverse, padding_predicate, pad_value); Block new_scope_block = Downcast(new_stmt); // Step 4: Rewrite buffer_map of the PrimFunc if necessary. @@ -1472,20 +1517,21 @@ struct TransformLayoutTraits : public UnpackedInstTraits private: static constexpr size_t kNumInputs = 2; - static constexpr size_t kNumAttrs = 3; + static constexpr size_t kNumAttrs = 4; static constexpr size_t kNumDecisions = 0; static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv, IndexMap index_map, Integer buffer_index, Integer buffer_index_type, - Optional pad_value) { + Optional pad_value, + Bool assume_injective_transform) { return sch->TransformLayout(block_rv, buffer_index.IntValue(), static_cast(buffer_index_type->value), index_map, - pad_value); + pad_value, assume_injective_transform.operator bool()); } static String UnpackedAsPython(Array outputs, String block_rv, IndexMap index_map, Integer buffer_index, Integer buffer_index_type, - Optional pad_value) { + Optional pad_value, Bool assume_injective_transform) { PythonAPICall py("transform_layout"); py.Input("block", block_rv); @@ -1495,6 +1541,7 @@ struct TransformLayoutTraits : public UnpackedInstTraits py.Input("buffer", os.str()); py.Input("index_map", index_map->ToPythonString()); py.Input("pad_value", pad_value ? pad_value.value()->ToPythonString() : "None"); + py.Input("assume_injective_transform", assume_injective_transform.operator bool()); return py.Str(); } @@ -1510,6 +1557,7 @@ struct TransformLayoutTraits : public UnpackedInstTraits } else { attrs_record.push_back(attrs[2]); } + attrs_record.push_back(attrs[3]); return std::move(attrs_record); } @@ -1523,6 +1571,7 @@ struct TransformLayoutTraits : public UnpackedInstTraits } else { attrs.push_back(attrs_record[2]); } + attrs.push_back(attrs_record[3]); return attrs; } diff --git a/src/tir/schedule/schedule.cc b/src/tir/schedule/schedule.cc index d008f3639c78..4177d916486b 100644 --- a/src/tir/schedule/schedule.cc +++ b/src/tir/schedule/schedule.cc @@ -253,10 +253,10 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnannotate") TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformLayout") .set_body_typed([](Schedule self, const BlockRV& block_rv, int buffer_index, int buffer_index_type, const IndexMap& index_map, - const Optional& pad_value) { + const Optional& pad_value, bool assume_injective_transform) { return self->TransformLayout(block_rv, buffer_index, static_cast(buffer_index_type), index_map, - pad_value); + pad_value, assume_injective_transform); }); TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformBlockLayout") .set_body_method(&ScheduleNode::TransformBlockLayout); diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc index 8852fccf881a..dba34c2ca3f3 100644 --- a/src/tir/schedule/traced_schedule.cc +++ b/src/tir/schedule/traced_schedule.cc @@ -523,15 +523,18 @@ void TracedScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann_k void TracedScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map, - const Optional& pad_value) { + const Optional& pad_value, + bool assume_injective_transform) { ConcreteScheduleNode::TransformLayout(block_rv, buffer_index, buffer_index_type, index_map, - pad_value); + pad_value, assume_injective_transform); static const InstructionKind& kind = InstructionKind::Get("TransformLayout"); trace_->Append( /*inst=*/Instruction( /*kind=*/kind, /*inputs=*/{block_rv, index_map}, - /*attrs=*/{Integer(buffer_index), Integer(buffer_index_type), pad_value}, + /*attrs=*/ + {Integer(buffer_index), Integer(buffer_index_type), pad_value, + Bool(assume_injective_transform)}, /*outputs=*/{})); } diff --git a/src/tir/schedule/traced_schedule.h b/src/tir/schedule/traced_schedule.h index ee65c721ad9f..7bd83855557d 100644 --- a/src/tir/schedule/traced_schedule.h +++ b/src/tir/schedule/traced_schedule.h @@ -107,7 +107,8 @@ class TracedScheduleNode : public ConcreteScheduleNode { void Unannotate(const BlockRV& block_rv, const String& ann_key) override; /******** Schedule: Layout transformation ********/ void TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, - const IndexMap& index_map, const Optional& pad_value) override; + const IndexMap& index_map, const Optional& pad_value, + bool assume_injective_transform) override; void TransformBlockLayout(const BlockRV& block_rv, const IndexMap& index_map) override; void SetAxisSeparator(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index 1561d962b3c8..c9a8f70ef7b3 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -477,11 +477,19 @@ class BasePaddingCompare(tvm.testing.CompareBeforeAfter): index_map = tvm.testing.parameter(lambda i: [i // 4, i % 4]) + assume_injective_transform = tvm.testing.parameter(False) + @pytest.fixture - def transform(self, pad_value, transformed_buffer, index_map): + def transform(self, pad_value, transformed_buffer, index_map, assume_injective_transform): def transform(mod): sch = tir.Schedule(mod) - sch.transform_layout("block", transformed_buffer, index_map, pad_value=pad_value) + sch.transform_layout( + "block", + transformed_buffer, + index_map, + pad_value=pad_value, + assume_injective_transform=assume_injective_transform, + ) return sch.mod return transform @@ -565,11 +573,26 @@ def expected(): A[i // 4, i % 4] = 0 -class TestImplicitPadding(BasePaddingCompare): - """When pad_value is None, the buffer can be implicitly padded. The padded region is not - accessed because the original loop extent is not changed. +class TestErrorIfPaddingForbidden(BasePaddingCompare): + """Unless padding is explicitly enabled, should raise error""" + + def before(): + A = T.alloc_buffer(14, "int32") + for i in T.serial(14): + with T.block("block"): + vi = T.axis.remap("S", [i]) + A[vi] = 0 + + expected = tvm.tir.schedule.schedule.ScheduleError + + +class TestImplicitPaddingAssumeInjective(BasePaddingCompare): + """When pad_value is None and assume_injective_transform is set, the buffer can be implicitly + padded. The padded region is not accessed because the original loop extent is not changed. """ + assume_injective_transform = tvm.testing.parameter(True) + def before(): A = T.alloc_buffer(14, "int32") for i in T.serial(14): From 3e0e1acfd6188100b9b7d421f42dc30c0b663dcb Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Tue, 28 Feb 2023 19:23:33 -0800 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: Siyuan Feng --- include/tvm/tir/schedule/schedule.h | 2 +- python/tvm/tir/schedule/schedule.py | 2 +- src/tir/schedule/primitive.h | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index cdc22a294d00..7f2bdf6b4ebb 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -643,7 +643,7 @@ class ScheduleNode : public runtime::Object { * optimizations may assume that this precondition is met, and * may result in incorrect results being returned. * - * \param assume_injective_transform If set to true, the schedule primitive will assume the + * \param assume_injective_transform If set to true, the schedule primitive will assume the * index_map is injective and skip checking overlapping of the mapped indices. This can be useful * for complicated index_map that the analysis does not cover. It is the callers' responsibility * to ensure the index map is injective, otherwise, the correctness of the schedule is not diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 5458f6e8bab6..87bb09f5370a 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -2645,7 +2645,7 @@ def transform_layout( index_map: Union[IndexMap, Callable], pad_value: Optional[Union[int, float, PrimExpr, IndexMap, Callable]] = None, *, - assume_injective_transform=False, + assume_injective_transform: bool = False, ) -> None: """Apply a transformation represented by IndexMap to buffer diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index 65e6834cac81..0b7a4f6280db 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -501,7 +501,7 @@ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const String& * \param buffer_index_type The type of the buffer index, kRead or kWrite. * \param index_map The transformation to apply. * \param pad_value The value to write into padding introduced by the transformation. - * \param assume_injective_transform If set to true, the schedule primitive will assume the + * \param assume_injective_transform If set to true, the schedule primitive will assume the * index_map is injective and skip checking overlapping of the mapped indices. This can be useful * for complicated index_map that the analysis does not cover. It is the callers' responsibility * to ensure the index map is injective, otherwise, the correctness of the schedule is not