From 50e1e92e5a20808fb40fcc0dce445d7aefdc15c9 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Tue, 8 Mar 2022 18:09:34 -0500 Subject: [PATCH 01/11] [TIR][Schedule] Transform layout --- include/tvm/tir/index_map.h | 16 ++ include/tvm/tir/schedule/schedule.h | 16 ++ python/tvm/tir/function.py | 40 ++- python/tvm/tir/schedule/schedule.py | 68 ++++- src/tir/ir/index_map.cc | 51 ++++ src/tir/schedule/analysis.h | 10 + src/tir/schedule/analysis/analysis.cc | 31 +++ src/tir/schedule/concrete_schedule.cc | 9 + src/tir/schedule/concrete_schedule.h | 4 +- src/tir/schedule/primitive.h | 16 ++ src/tir/schedule/primitive/block_annotate.cc | 38 --- .../primitive/layout_transformation.cc | 239 ++++++++++++++++++ src/tir/schedule/schedule.cc | 4 + src/tir/schedule/traced_schedule.cc | 13 + src/tir/schedule/traced_schedule.h | 3 + .../test_tir_schedule_transform_layout.py | 160 ++++++++++++ 16 files changed, 677 insertions(+), 41 deletions(-) create mode 100644 src/tir/schedule/primitive/layout_transformation.cc create mode 100644 tests/python/unittest/test_tir_schedule_transform_layout.py diff --git a/include/tvm/tir/index_map.h b/include/tvm/tir/index_map.h index 237111306c2a..195bf7e02ce3 100644 --- a/include/tvm/tir/index_map.h +++ b/include/tvm/tir/index_map.h @@ -106,11 +106,19 @@ class IndexMapNode : public Object { */ Array MapShape(const Array& shape) const; + /*! + * \brief Convert to string representation in Python. + * \return The stringified lambda expression in Python. + */ + String ToPythonString() const; + void VisitAttrs(AttrVisitor* v) { v->Visit("initial_indices", &initial_indices); v->Visit("final_indices", &final_indices); } + static constexpr const char* _type_key = "tir.IndexMap"; + TVM_DECLARE_FINAL_OBJECT_INFO(IndexMapNode, Object); }; @@ -118,6 +126,14 @@ class IndexMap : public ObjectRef { public: IndexMap(Array initial_indices, Array final_indices); + /*! + * \brief Create an index map from a packed function + * \param ndim The number of dimensions + * \param func The function to be applied + * \return The created index map + */ + static IndexMap FromFunc(int ndim, runtime::TypedPackedFunc(Array)> func); + /*! \brief Generate the inverse mapping. * * The range of the input indices is required in order to ensure diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index be06b44820cd..1507955f5cfc 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -20,6 +20,7 @@ #define TVM_TIR_SCHEDULE_SCHEDULE_H_ #include +#include #include #include @@ -521,6 +522,21 @@ class ScheduleNode : public runtime::Object { */ virtual void Unannotate(const BlockRV& block_rv, const String& ann_key) = 0; + /******** Schedule: Layout transformation ********/ + /*! + * \brief Apply a transformation represented by IndexMap to buffer + * \details The indices and the access region to the target buffer is transformed by the given + * index_map. The index_map is used to infer the new shape of the buffer. Buffer must be either + * a function parameter, or allocated in a block (it cannot be a buffer subregion created via + * 'match_buffer'). + * \param block_rv The block that accesses the target buffer. + * \param buffer_index The index of the buffer in block's read or write region. + * \param is_write_index Whether the buffer_index is the index of the block's write region. + * \param index_map The transformation to apply. + */ + virtual void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, + const IndexMap& index_map) = 0; + /******** Schedule: Misc ********/ /*! \brief A no-op that marks the start of postprocessing phase of scheduling */ virtual void EnterPostproc() = 0; diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py index fdee18f88cf8..56d03073bb6c 100644 --- a/python/tvm/tir/function.py +++ b/python/tvm/tir/function.py @@ -16,7 +16,8 @@ # under the License. """Function data types.""" -from typing import Mapping, Union +from typing import Callable, List, Mapping, Union +import inspect import tvm._ffi import tvm.runtime @@ -239,3 +240,40 @@ def get(name: str): The TensorIntrin with the specified name. """ return _ffi_api.TensorIntrinGet(name) # pylint: type: ignore + + +@tvm._ffi.register_object("tir.IndexMap") +class IndexMap(Object): + """A mapping from multi-dimensional indices to another set of multi-dimensional indices + + Parameters + ---------- + initial_indices : List[Var] + Variables representing the indices prior to remapping. + final_indices : List[PrimExpr] + Expressions defining the indices after remapping. + """ + + initial_indices: List[Var] + final_indices: List[PrimExpr] + + @staticmethod + def from_func(func: Callable) -> "IndexMap": + """Create an index map from a function + + Parameters + ---------- + func : Callable + The function to map from source indices to target indices + """ + + def wrap(args: List[Var]) -> List[PrimExpr]: + result = func(*args) + if isinstance(result, tuple): + return list(result) + if not isinstance(result, list): + result = [result] + return result + + ndim = len(inspect.signature(func).parameters) + return _ffi_api.IndexMapFromFunc(ndim, wrap) # type: ignore # pylint: disable=no-member diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 96fa21f30020..9849554e813f 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -15,13 +15,14 @@ # specific language governing permissions and limitations # under the License. """The TensorIR schedule class""" -from typing import Dict, List, Optional, Union +from typing import Callable, Dict, List, Optional, Union from tvm._ffi import register_object as _register_object from tvm.error import TVMError, register_error from tvm.ir import IRModule, PrimExpr from tvm.runtime import Object, String from tvm.tir import Block, FloatImm, For, IntImm, PrimFunc +from ..function import IndexMap from . import _ffi_api from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod @@ -2111,6 +2112,71 @@ def after_unannotate(a: T.handle, b: T.handle) -> None: self, block_or_loop, ann_key ) + ########## Schedule: Layout transformation ########## + + def transform_layout( + self, + block: BlockRV, + buffer_index: int, + is_write_index: bool, + index_map: Union[IndexMap, Callable], + ) -> None: + """Apply a transformation represented by IndexMap to buffer + Parameters + ---------- + block_rv : BlockRV + The block that accesses the target buffer + buffer_index: int + The index of the buffer in block's read or write region + is_write_index : bool + Whether the buffer_index is the index of the block's write region + index_map : Union[IndexMap, Callable] + The transformation to apply + Examples + -------- + Before transform_layout, in TensorIR, the IR is: + .. code-block:: python + @T.prim_func + def before_transform_layout(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128), "float32") + B = T.alloc_buffer((128, 128), "float32") + C = T.match_buffer(c, (128, 128), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = B[vi, vj] + 1.0 + Create the schedule and do transform_layout: + .. code-block:: python + sch = tir.Schedule(before_storage_align) + sch.transform_layout(sch.get_block("B"), buffer_index=0, is_write_index=True, + index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16)) + print(sch.mod["main"].script()) + After applying transform_layout, the IR becomes: + .. code-block:: python + @T.prim_func + def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128), "float32") + B = T.alloc_buffer((8, 8, 16, 16), "float32") + C = T.match_buffer(c, (128, 128), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0 + """ + if callable(index_map): + index_map = IndexMap.from_func(index_map) + _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member + self, block, buffer_index, is_write_index, index_map + ) + ########## Schedule: Misc ########## @type_checked diff --git a/src/tir/ir/index_map.cc b/src/tir/ir/index_map.cc index ba0998e84ffc..7955c4f7fc6f 100644 --- a/src/tir/ir/index_map.cc +++ b/src/tir/ir/index_map.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include @@ -40,6 +41,15 @@ IndexMap::IndexMap(Array initial_indices, Array final_indices) { data_ = std::move(n); } +IndexMap IndexMap::FromFunc(int ndim, runtime::TypedPackedFunc(Array)> func) { + Array initial_indices; + initial_indices.reserve(ndim); + for (int i = 0; i < ndim; ++i) { + initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32))); + } + return IndexMap(initial_indices, func(initial_indices)); +} + IndexMap IndexMap::Inverse(Array initial_ranges) const { // Dummy variables to represent the inverse's inputs. Array output_vars; @@ -142,6 +152,40 @@ Array IndexMapNode::MapShape(const Array& shape) const { return output; } +String IndexMapNode::ToPythonString() const { + std::unordered_set used_names; + Map var_remap; + for (const Var& initial_index : initial_indices) { + if (used_names.count(initial_index->name_hint)) { + std::string new_name = initial_index->name_hint + std::to_string(used_names.size()); + used_names.insert(new_name); + var_remap.Set(initial_index, Var(new_name)); + } else { + used_names.insert(initial_index->name_hint); + } + } + std::ostringstream oss; + oss << "lambda "; + for (size_t i = 0; i < initial_indices.size(); ++i) { + if (i != 0) { + oss << ", "; + } + auto it = var_remap.find(initial_indices[i]); + if (it != var_remap.end()) { + oss << (*it).second; + } else { + oss << initial_indices[i]; + } + } + oss << ": ("; + for (size_t i = 0; i < final_indices.size(); ++i) { + oss << Substitute(final_indices[i], var_remap); + oss << ", "; + } + oss << ")"; + return String(oss.str()); +} + TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) .set_dispatch([](const ObjectRef& node, ReprPrinter* p) { auto* op = static_cast(node.get()); @@ -150,5 +194,12 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) TVM_REGISTER_NODE_TYPE(IndexMapNode); +TVM_REGISTER_GLOBAL("tir.IndexMap") + .set_body_typed([](Array initial_indices, Array final_indices) { + return IndexMap(initial_indices, final_indices); + }); + +TVM_REGISTER_GLOBAL("tir.IndexMapFromFunc").set_body_typed(IndexMap::FromFunc); + } // namespace tir } // namespace tvm diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h index 9c6d1e6e96da..9414bb1734ef 100644 --- a/src/tir/schedule/analysis.h +++ b/src/tir/schedule/analysis.h @@ -401,6 +401,16 @@ struct ProducerConsumerSplit { */ Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, bool is_write); +/*! + * \brief Find the defining site of the buffer in the given block and its ancestors + * \param block_sref The block sref + * \param buffer The buffer + * \return The defining site of the buffer and whether the buffer is allocated (otherwise the + * buffer is from match_buffer). + */ +std::pair, bool> GetBufferDefiningSite(const StmtSRef& block_sref, + const Buffer& buffer); + /******** Reduction Block Related ********/ /*! diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc index c7ed67187793..59a805fbdf2d 100644 --- a/src/tir/schedule/analysis/analysis.cc +++ b/src/tir/schedule/analysis/analysis.cc @@ -1029,6 +1029,37 @@ Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n, return access_region[n]->buffer; } +std::pair, bool> GetBufferDefiningSite(const StmtSRef& block_sref, + const Buffer& buffer) { + // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or + // match_buffers. + const StmtSRefNode* defining_site_sref = block_sref.get(); + while (defining_site_sref != nullptr) { + const auto* block = defining_site_sref->StmtAs(); + // If this sref is not a block sref, skip it. + if (block == nullptr) { + defining_site_sref = defining_site_sref->parent; + continue; + } + // Try to find the buffer in `allloc_buffers` + for (const Buffer& alloc_buffer : block->alloc_buffers) { + if (buffer.same_as(alloc_buffer)) { + return {GetRef(defining_site_sref), true}; + } + } + // We do not allow the buffer being defined in `match_buffer`. + for (const MatchBufferRegion match_buffer : block->match_buffers) { + if (buffer.same_as(match_buffer)) { + return {GetRef(defining_site_sref), false}; + } + } + defining_site_sref = defining_site_sref->parent; + } + // If we cannot find the defining site block, it means that the buffer must be in the function's + // buffer_map, which isn't an intermediate buffer. + return {NullOpt, false}; +} + /******** Pattern Matcher ********/ /*! diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index 394f0f26db35..e1cf46bfe6ec 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -685,6 +685,15 @@ void ConcreteScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann TVM_TIR_SCHEDULE_END("unannotate", this->error_render_level_); } +/******** Schedule: Layout transformation ********/ +void ConcreteScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, + bool is_write_index, const IndexMap& index_map) { + TVM_TIR_SCHEDULE_BEGIN(); + tir::TransformLayout(state_, this->GetSRef(block_rv), buffer_index, is_write_index, index_map); + this->state_->DebugVerify(); + TVM_TIR_SCHEDULE_END("transform_layout", this->error_render_level_); +} + /******** Schedule: Misc ********/ } // namespace tir diff --git a/src/tir/schedule/concrete_schedule.h b/src/tir/schedule/concrete_schedule.h index f0f25ecafa3a..f695d523b46c 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -131,7 +131,9 @@ class ConcreteScheduleNode : public ScheduleNode { void Unannotate(const LoopRV& loop_rv, const String& ann_key) override; void Annotate(const BlockRV& block_rv, const String& ann_key, const ObjectRef& ann_val) override; void Unannotate(const BlockRV& block_rv, const String& ann_key) override; - + /******** Schedule: Layout transformation ********/ + void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, + const IndexMap& index_map) override; /******** Schedule: Misc ********/ void EnterPostproc() override {} diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index 0cd2d3e6f38a..b32b99fb446b 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -415,6 +415,22 @@ TVM_DLL void Annotate(ScheduleState self, const StmtSRef& sref, const String& an */ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const String& ann_key); +/******** Schedule: Layout transformation ********/ +/*! + * \brief Apply a transformation represented by IndexMap to buffer + * \details The indices and the access region to the target buffer is transformed by the given + * index_map. The index_map is also used to infer the new shape of the buffer. Buffer must be + * one of the parameter of the function, or allocated in some blocks (it cannot be a buffer + * subregion created via match_buffer). + * \param self The state of the schedule + * \param block_sref The block sref that accesses the target buffer. + * \param buffer_index The index of the buffer in block's read or write region. + * \param is_write_index Whether the buffer_index is the index of the block's write region. + * \param index_map The transformation to apply. + */ +TVM_DLL void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, + bool is_write_index, const IndexMap& index_map); + /******** Schedule: Misc ********/ } // namespace tir diff --git a/src/tir/schedule/primitive/block_annotate.cc b/src/tir/schedule/primitive/block_annotate.cc index 418e770a5c93..f9cec421cd21 100644 --- a/src/tir/schedule/primitive/block_annotate.cc +++ b/src/tir/schedule/primitive/block_annotate.cc @@ -64,44 +64,6 @@ class StorageAlignAxisOutOfRangeError : public ScheduleError { int axis_; }; -/*! - * \brief Find the defining site of the buffer in the given block and its ancestors - * \param block_sref The block sref - * \param buffer The buffer - * \return The defining site of the buffer and whether the buffer is allocated (otherwise the - * buffer is from match_buffer). - */ -std::pair, bool> GetBufferDefiningSite(const StmtSRef& block_sref, - const Buffer& buffer) { - // Climb up along the sref tree, and find the block where `buffer` is in alloc_buffers or - // match_buffers. - const StmtSRefNode* defining_site_sref = block_sref.get(); - while (defining_site_sref != nullptr) { - const auto* block = defining_site_sref->StmtAs(); - // If this sref is not a block sref, skip it. - if (block == nullptr) { - defining_site_sref = defining_site_sref->parent; - continue; - } - // Try to find the buffer in `allloc_buffers` - for (const Buffer& alloc_buffer : block->alloc_buffers) { - if (buffer.same_as(alloc_buffer)) { - return {GetRef(defining_site_sref), true}; - } - } - // We do not allow the buffer being defined in `match_buffer`. - for (const MatchBufferRegion match_buffer : block->match_buffers) { - if (buffer.same_as(match_buffer)) { - return {GetRef(defining_site_sref), false}; - } - } - defining_site_sref = defining_site_sref->parent; - } - // If we cannot find the defining site block, it means that the buffer must be in the function's - // buffer_map, which isn't an intermediate buffer. - return {NullOpt, false}; -} - class NonAllocatedBufferError : public ScheduleError { public: explicit NonAllocatedBufferError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {} diff --git a/src/tir/schedule/primitive/layout_transformation.cc b/src/tir/schedule/primitive/layout_transformation.cc new file mode 100644 index 000000000000..4570ab43ee9a --- /dev/null +++ b/src/tir/schedule/primitive/layout_transformation.cc @@ -0,0 +1,239 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include "../utils.h" + +namespace tvm { +namespace tir { + +class TransformLayoutRewriter : private StmtExprMutator { + public: + /*! + * \brief Rewrite the access to the buffer after the transformation + * \param scope_stmt The parent statement that contains all accesses to the target buffer + * \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 + * \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 Stmt& scope_stmt, + const Buffer& old_buffer, + const Buffer& new_buffer, + const IndexMap& index_map) { + TransformLayoutRewriter rewriter(old_buffer, new_buffer, index_map); + Stmt result = rewriter(scope_stmt); + return {result, rewriter.block_sref_reuse_}; + } + + private: + TransformLayoutRewriter(const Buffer& old_buffer, const Buffer& new_buffer, + const IndexMap& index_map) + : old_buffer_(old_buffer), + new_buffer_(new_buffer), + index_map_(index_map), + buffer_data_to_buffer_{{new_buffer->data, new_buffer}} {} + + void RewriteBufferAccess(Buffer* buffer, Array* indices) { + *buffer = new_buffer_; + *indices = index_map_->MapIndices(*indices); + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + BufferLoad buffer_load = Downcast(StmtExprMutator::VisitExpr_(op)); + if (buffer_load->buffer.same_as(old_buffer_)) { + auto* n = buffer_load.CopyOnWrite(); + RewriteBufferAccess(&n->buffer, &n->indices); + } + return std::move(buffer_load); + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + BufferStore buffer_store = Downcast(StmtExprMutator::VisitStmt_(op)); + if (buffer_store->buffer.same_as(old_buffer_)) { + auto* n = buffer_store.CopyOnWrite(); + RewriteBufferAccess(&n->buffer, &n->indices); + } + return std::move(buffer_store); + } + + void RewriteAccessRegion(Array* old_access_regions, + const Array& infered_access_regions) { + auto fmutate = [this, &infered_access_regions](const BufferRegion& buffer_region) { + if (buffer_region->buffer.same_as(old_buffer_)) { + ICHECK(infered_access_regions.size() == 1); + return infered_access_regions[0]; + } + return buffer_region; + }; + (*old_access_regions).MutateByApply(fmutate); + } + + Stmt VisitStmt_(const BlockNode* op) final { + Block block = Downcast(StmtExprMutator::VisitStmt_(op)); + auto infered_access_regions = GetBlockReadWriteRegion(block, buffer_data_to_buffer_); + auto* n = block.CopyOnWrite(); + RewriteAccessRegion(&n->reads, infered_access_regions[0]); + RewriteAccessRegion(&n->writes, infered_access_regions[1]); + block_sref_reuse_.Set(GetRef(op), block); + return std::move(block); + } + + const Buffer& old_buffer_; + const Buffer& new_buffer_; + const IndexMap& index_map_; + Map buffer_data_to_buffer_; + Map block_sref_reuse_; +}; + +class BufferIsSubregionError : public ScheduleError { + public: + explicit BufferIsSubregionError(IRModule mod, Buffer buffer) : mod_(mod), buffer_(buffer) {} + + String FastErrorString() const final { + return "ScheduleError: The input buffer is defined in `match_buffer` of a block, it is expected" + " to be a function parameter or allocated by a block"; + } + + String DetailRenderTemplate() const final { + std::ostringstream os; + os << "ScheduleError: The input buffer " << buffer_->name << " is defined in `match_buffer` of " + << "a block, it is expected to be a function parameter or allocated by a block."; + return os.str(); + } + + Array LocationsOfInterest() const final { return {}; } + IRModule mod() const final { return mod_; } + + private: + IRModule mod_; + Buffer buffer_; +}; + +void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, + bool is_write_index, const IndexMap& index_map) { + const BlockNode* block_ptr = TVM_SREF_TO_BLOCK(block_ptr, block_sref); + Buffer old_buffer = GetNthAccessBuffer(self, GetRef(block_ptr), buffer_index, + /*is_write=*/is_write_index); + Optional defining_site_sref; + bool is_alloc; + std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, old_buffer); + if (defining_site_sref.defined() && !is_alloc) { + throw BufferIsSubregionError(self->mod, old_buffer); + } + + StmtSRef scope_sref = defining_site_sref.defined() + ? defining_site_sref.value() + : GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false); + const BlockNode* scope_block = TVM_SREF_TO_BLOCK(scope_block, scope_sref); + + // Step 1: Infer the shape of the new buffer + ObjectPtr new_buffer_node = make_object(*(old_buffer.get())); + new_buffer_node->shape = index_map->MapShape(old_buffer->shape); + Buffer new_buffer{new_buffer_node}; + + // Step 2: Rewrite access indices and regions of the buffer + Stmt new_stmt; + Map block_sref_reuse; + std::tie(new_stmt, block_sref_reuse) = TransformLayoutRewriter::Rewrite( + GetRef(scope_block), old_buffer, new_buffer, index_map); + Block new_scope_block = Downcast(new_stmt); + + // Step 3: Rewrite alloc_buffer of the block or buffer_map of the PrimFunc. + if (defining_site_sref.defined()) { + auto* n = new_scope_block.CopyOnWrite(); + n->alloc_buffers.MutateByApply([&old_buffer, &new_buffer](const Buffer& buffer) { + if (buffer.same_as(old_buffer)) { + return new_buffer; + } + return buffer; + }); + block_sref_reuse.Set(GetRef(scope_block), new_scope_block); + } else { + GlobalVar g_var; + GetRootPrimFunc(self->mod, scope_block, &g_var); + IRModuleNode* new_mod = self->mod.CopyOnWrite(); + MapNode* new_map = new_mod->functions.CopyOnWrite(); + PrimFunc ref_new_func = Downcast(std::move(new_map->at(g_var))); + PrimFuncNode* new_func = ref_new_func.CopyOnWrite(); + MapNode* new_buffer_map = new_func->buffer_map.CopyOnWrite(); + for (auto it = new_buffer_map->begin(); it != new_buffer_map->end(); ++it) { + if ((*it).second.same_as(old_buffer)) { + (*it).second = new_buffer; + } + } + new_map->at(g_var) = std::move(ref_new_func); + } + + // Step 4: Replace the scope block with the new block + self->Replace(scope_sref, new_scope_block, block_sref_reuse); +} + +/******** InstructionKind Registration ********/ + +struct TransformLayoutTraits : public UnpackedInstTraits { + static constexpr const char* kName = "TransformLayout"; + static constexpr bool kIsPure = false; + + private: + static constexpr size_t kNumInputs = 1; + static constexpr size_t kNumAttrs = 3; + static constexpr size_t kNumDecisions = 0; + + static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv, Integer buffer_index, + Bool is_write_index, IndexMap index_map) { + return sch->TransformLayout(block_rv, buffer_index, is_write_index, index_map); + } + + static String UnpackedAsPython(Array outputs, String block_rv, Integer buffer_index, + Bool is_write_index, IndexMap index_map) { + PythonAPICall py("transform_layout"); + py.Input("block", block_rv); + py.Input("buffer_index", buffer_index); + py.Input("is_write_index", is_write_index); + py.Input("index_map", index_map->ToPythonString()); + return py.Str(); + } + + public: + static ObjectRef AttrsAsJSON(const Array& attrs) { + Array attrs_record; + attrs_record.reserve(kNumAttrs); + attrs_record.push_back(attrs[0]); + attrs_record.push_back(attrs[1]); + attrs_record.push_back(String(::tvm::SaveJSON(attrs[2]))); + return std::move(attrs_record); + } + + static Array AttrsFromJSON(const ObjectRef& attrs_record_) { + Array attrs_record = Downcast>(attrs_record_); + Array attrs; + attrs.push_back(attrs_record[0]); + attrs.push_back(attrs_record[1]); + attrs.push_back(::tvm::LoadJSON(Downcast(attrs_record[2]))); + return attrs; + } + + template + friend struct ::tvm::tir::UnpackedInstTraits; +}; + +TVM_REGISTER_INST_KIND_TRAITS(TransformLayoutTraits); + +} // namespace tir +} // namespace tvm diff --git a/src/tir/schedule/schedule.cc b/src/tir/schedule/schedule.cc index b466843f9459..060f9fbe7619 100644 --- a/src/tir/schedule/schedule.cc +++ b/src/tir/schedule/schedule.cc @@ -226,6 +226,10 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnannotate") throw; }); +/******** (FFI) Layout transformation ********/ +TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformLayout") + .set_body_method(&ScheduleNode::TransformLayout); + /******** (FFI) Misc ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleEnterPostproc") .set_body_method(&ScheduleNode::EnterPostproc); diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc index 1e2e57eb6eca..1d41d9e443eb 100644 --- a/src/tir/schedule/traced_schedule.cc +++ b/src/tir/schedule/traced_schedule.cc @@ -427,6 +427,19 @@ void TracedScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann_k /*outputs=*/{})); } +/******** Schedule: Layout transformation ********/ + +void TracedScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, + bool is_write_index, const IndexMap& index_map) { + ConcreteScheduleNode::TransformLayout(block_rv, buffer_index, is_write_index, index_map); + static const InstructionKind& kind = InstructionKind::Get("TransformLayout"); + trace_->Append( + /*inst=*/Instruction(/*kind=*/kind, + /*inputs=*/{block_rv}, + /*attrs=*/{Integer(buffer_index), Bool(is_write_index), index_map}, + /*outputs=*/{})); +} + /******** Schedule: Misc ********/ void TracedScheduleNode::EnterPostproc() { diff --git a/src/tir/schedule/traced_schedule.h b/src/tir/schedule/traced_schedule.h index 5d3fdbf570de..9ad9b237bc1a 100644 --- a/src/tir/schedule/traced_schedule.h +++ b/src/tir/schedule/traced_schedule.h @@ -95,6 +95,9 @@ class TracedScheduleNode : public ConcreteScheduleNode { void Unannotate(const LoopRV& loop_rv, const String& ann_key) override; void Annotate(const BlockRV& block_rv, const String& ann_key, const ObjectRef& ann_val) override; void Unannotate(const BlockRV& block_rv, const String& ann_key) override; + /******** Schedule: Layout transformation ********/ + void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, + const IndexMap& index_map) override; /******** Schedule: Misc ********/ void EnterPostproc() final; }; diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py new file mode 100644 index 000000000000..f08441503e19 --- /dev/null +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -0,0 +1,160 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=missing-function-docstring,missing-module-docstring +import sys + +import pytest + +import tvm +from tvm import tir +from tvm.script import tir as T +from tvm.tir.schedule.testing import verify_trace_roundtrip + +# fmt: off +# pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks + +def packed_index_map_func(m, n): + return m // 16, n // 16, m % 16, n % 16 + + +@T.prim_func +def two_elementwise(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128), "float32") + B = T.alloc_buffer((128, 128), "float32") + C = T.match_buffer(c, (128, 128), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = B[vi, vj] + 1.0 + + +@T.prim_func +def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128), "float32") + B = T.alloc_buffer((8, 8, 16, 16), "float32") + C = T.match_buffer(c, (128, 128), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi // 16, vj // 16, vi % 16, vj % 16] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0 + + +@T.prim_func +def two_elementwise_transformed_input_buffer(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (8, 8, 16, 16), "float32") + B = T.alloc_buffer((128, 128), "float32") + C = T.match_buffer(c, (128, 128), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi // 16, vj // 16, vi % 16, vj % 16] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = B[vi, vj] + 1.0 + + +@T.prim_func +def two_elementwise_transformed_output_buffer(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128), "float32") + B = T.alloc_buffer((128, 128), "float32") + C = T.match_buffer(c, (8, 8, 16, 16), "float32") + for i, j in T.grid(128, 128): + with T.block("B"): + vi, vj = T.axis.remap("SS", [i, j]) + B[vi, vj] = A[vi, vj] * 2.0 + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi // 16, vj // 16, vi % 16, vj % 16] = B[vi, vj] + 1.0 + + +@T.prim_func +def permuted_shared_memory(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128)) + C = T.match_buffer(c, (128, 128)) + A_shared = T.alloc_buffer((128, 128), scope="shared") + for i0, j0, in T.grid(32, 4): + for fused_i1_j1 in T.thread_binding(0, 32, 'threadIdx.x'): + for j2 in T.vectorized(0, 4): + with T.block("A_shared"): + vi = T.axis.S(128, i0 * 4 + fused_i1_j1 // 8) + vj = T.axis.S(128, j0 * 32 + fused_i1_j1 % 8 * 4 + j2) + A_shared[vi, vj] = A[vi, vj] + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = A_shared[vi, vj] + 1.0 + + +@T.prim_func +def permuted_shared_memory_transformed(a: T.handle, c: T.handle) -> None: + A = T.match_buffer(a, (128, 128)) + C = T.match_buffer(c, (128, 128)) + A_shared = T.alloc_buffer((32, 4, 4, 32), scope="shared") + for i0, j0, in T.grid(32, 4): + for fused_i1_j1 in T.thread_binding(0, 32, 'threadIdx.x'): + for j2 in T.vectorized(0, 4): + with T.block("A_shared"): + vi = T.axis.S(128, i0 * 4 + fused_i1_j1 // 8) + vj = T.axis.S(128, j0 * 32 + fused_i1_j1 % 8 * 4 + j2) + A_shared[vi // 4, vj // 32, vi % 4, (((vj % 32) // 8) ^ (vi % 4)) + vj % 8] = A[vi, vj] + for i, j in T.grid(128, 128): + with T.block("C"): + vi, vj = T.axis.remap("SS", [i, j]) + C[vi, vj] = A_shared[vi // 4, vj // 32, vi % 4, (((vj % 32) // 8) ^ (vi % 4)) + vj % 8] + 1.0 + + +# pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks +# fmt: on + + +def test_two_elementwise_transform_intermediate_buffer(): + sch = tir.Schedule(two_elementwise, debug_mask="all") + block = sch.get_block("B") + sch.transform_layout(block, 0, False, packed_index_map_func) + tvm.ir.assert_structural_equal(two_elementwise_transformed_intermediate_buffer, sch.mod["main"]) + verify_trace_roundtrip(sch=sch, mod=two_elementwise) + + +def test_two_elementwise_transform_input_buffer(): + sch = tir.Schedule(two_elementwise, debug_mask="all") + block = sch.get_block("B") + sch.transform_layout(block, 0, True, packed_index_map_func) + print(sch.mod["main"].script()) + tvm.ir.assert_structural_equal(two_elementwise_transformed_input_buffer, sch.mod["main"]) + verify_trace_roundtrip(sch=sch, mod=two_elementwise) + + +def test_two_elementwise_transform_output_buffer(): + sch = tir.Schedule(two_elementwise, debug_mask="all") + block = sch.get_block("C") + sch.transform_layout(block, 0, False, packed_index_map_func) + tvm.ir.assert_structural_equal(two_elementwise_transformed_output_buffer, sch.mod["main"]) + verify_trace_roundtrip(sch=sch, mod=two_elementwise) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__] + sys.argv[1:])) From ab3f83d0c7183ce29a69884991afc733c7cf76f1 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Wed, 9 Mar 2022 19:42:16 -0500 Subject: [PATCH 02/11] address commens --- python/tvm/tir/schedule/schedule.py | 2 + .../test_tir_schedule_transform_layout.py | 38 +------------------ 2 files changed, 3 insertions(+), 37 deletions(-) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 9849554e813f..e8431faacffc 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -2114,6 +2114,7 @@ def after_unannotate(a: T.handle, b: T.handle) -> None: ########## Schedule: Layout transformation ########## + @type_checked def transform_layout( self, block: BlockRV, @@ -2132,6 +2133,7 @@ def transform_layout( Whether the buffer_index is the index of the block's write region index_map : Union[IndexMap, Callable] The transformation to apply + Examples -------- Before transform_layout, in TensorIR, the IR is: diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index f08441503e19..ee92360c8228 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -91,42 +91,6 @@ def two_elementwise_transformed_output_buffer(a: T.handle, c: T.handle) -> None: C[vi // 16, vj // 16, vi % 16, vj % 16] = B[vi, vj] + 1.0 -@T.prim_func -def permuted_shared_memory(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128)) - C = T.match_buffer(c, (128, 128)) - A_shared = T.alloc_buffer((128, 128), scope="shared") - for i0, j0, in T.grid(32, 4): - for fused_i1_j1 in T.thread_binding(0, 32, 'threadIdx.x'): - for j2 in T.vectorized(0, 4): - with T.block("A_shared"): - vi = T.axis.S(128, i0 * 4 + fused_i1_j1 // 8) - vj = T.axis.S(128, j0 * 32 + fused_i1_j1 % 8 * 4 + j2) - A_shared[vi, vj] = A[vi, vj] - for i, j in T.grid(128, 128): - with T.block("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = A_shared[vi, vj] + 1.0 - - -@T.prim_func -def permuted_shared_memory_transformed(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128)) - C = T.match_buffer(c, (128, 128)) - A_shared = T.alloc_buffer((32, 4, 4, 32), scope="shared") - for i0, j0, in T.grid(32, 4): - for fused_i1_j1 in T.thread_binding(0, 32, 'threadIdx.x'): - for j2 in T.vectorized(0, 4): - with T.block("A_shared"): - vi = T.axis.S(128, i0 * 4 + fused_i1_j1 // 8) - vj = T.axis.S(128, j0 * 32 + fused_i1_j1 % 8 * 4 + j2) - A_shared[vi // 4, vj // 32, vi % 4, (((vj % 32) // 8) ^ (vi % 4)) + vj % 8] = A[vi, vj] - for i, j in T.grid(128, 128): - with T.block("C"): - vi, vj = T.axis.remap("SS", [i, j]) - C[vi, vj] = A_shared[vi // 4, vj // 32, vi % 4, (((vj % 32) // 8) ^ (vi % 4)) + vj % 8] + 1.0 - - # pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks # fmt: on @@ -134,7 +98,7 @@ def permuted_shared_memory_transformed(a: T.handle, c: T.handle) -> None: def test_two_elementwise_transform_intermediate_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, False, packed_index_map_func) + sch.transform_layout(block, 0, False, lambda m, n: m // 16, n // 16, m % 16, n % 16) tvm.ir.assert_structural_equal(two_elementwise_transformed_intermediate_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) From 203ea83e4dbb623290d06fe7275d46a2d5159e78 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Thu, 10 Mar 2022 12:47:09 -0500 Subject: [PATCH 03/11] fix --- .../python/unittest/test_tir_schedule_transform_layout.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index ee92360c8228..8d489fe601b4 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -98,7 +98,7 @@ def two_elementwise_transformed_output_buffer(a: T.handle, c: T.handle) -> None: def test_two_elementwise_transform_intermediate_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, False, lambda m, n: m // 16, n // 16, m % 16, n % 16) + sch.transform_layout(block, 0, True, lambda m, n: (m // 16, n // 16, m % 16, n % 16)) tvm.ir.assert_structural_equal(two_elementwise_transformed_intermediate_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) @@ -106,8 +106,7 @@ def test_two_elementwise_transform_intermediate_buffer(): def test_two_elementwise_transform_input_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, True, packed_index_map_func) - print(sch.mod["main"].script()) + sch.transform_layout(block, 0, False, packed_index_map_func) tvm.ir.assert_structural_equal(two_elementwise_transformed_input_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) @@ -115,7 +114,7 @@ def test_two_elementwise_transform_input_buffer(): def test_two_elementwise_transform_output_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("C") - sch.transform_layout(block, 0, False, packed_index_map_func) + sch.transform_layout(block, 0, True, packed_index_map_func) tvm.ir.assert_structural_equal(two_elementwise_transformed_output_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) From b4f58dfe3f9aa8f7f88da657d5737e0d86d69f62 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Fri, 11 Mar 2022 19:03:05 -0500 Subject: [PATCH 04/11] doc --- python/tvm/tir/schedule/schedule.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index e8431faacffc..622ca6c4c40d 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -2137,7 +2137,9 @@ def transform_layout( Examples -------- Before transform_layout, in TensorIR, the IR is: + .. code-block:: python + @T.prim_func def before_transform_layout(a: T.handle, c: T.handle) -> None: A = T.match_buffer(a, (128, 128), "float32") @@ -2151,14 +2153,20 @@ def before_transform_layout(a: T.handle, c: T.handle) -> None: with T.block("C"): vi, vj = T.axis.remap("SS", [i, j]) C[vi, vj] = B[vi, vj] + 1.0 + Create the schedule and do transform_layout: + .. code-block:: python + sch = tir.Schedule(before_storage_align) sch.transform_layout(sch.get_block("B"), buffer_index=0, is_write_index=True, index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16)) print(sch.mod["main"].script()) + After applying transform_layout, the IR becomes: + .. code-block:: python + @T.prim_func def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None: A = T.match_buffer(a, (128, 128), "float32") @@ -2172,6 +2180,7 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> with T.block("C"): vi, vj = T.axis.remap("SS", [i, j]) C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0 + """ if callable(index_map): index_map = IndexMap.from_func(index_map) From 1e0c44503bb27266c0a9ab464c6d0ee375f43687 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 14 Mar 2022 19:50:29 -0400 Subject: [PATCH 05/11] Address comments --- python/tvm/tir/function.py | 32 +++++++++++-------- src/tir/ir/index_map.cc | 2 +- .../test_tir_schedule_transform_layout.py | 23 +++++++------ 3 files changed, 31 insertions(+), 26 deletions(-) diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py index 56d03073bb6c..031ab7ce2c68 100644 --- a/python/tvm/tir/function.py +++ b/python/tvm/tir/function.py @@ -257,23 +257,29 @@ class IndexMap(Object): initial_indices: List[Var] final_indices: List[PrimExpr] + def __init__(self, initial_indices, final_indices): + self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices) + @staticmethod - def from_func(func: Callable) -> "IndexMap": + def from_func(mapping_function: Callable): """Create an index map from a function Parameters ---------- - func : Callable + mapping_function : Callable The function to map from source indices to target indices """ - - def wrap(args: List[Var]) -> List[PrimExpr]: - result = func(*args) - if isinstance(result, tuple): - return list(result) - if not isinstance(result, list): - result = [result] - return result - - ndim = len(inspect.signature(func).parameters) - return _ffi_api.IndexMapFromFunc(ndim, wrap) # type: ignore # pylint: disable=no-member + params = inspect.signature(mapping_function).parameters + default_index_dtype = "int32" + args = [] + for name, param in params.items(): + if param.kind in [ + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ]: + print(name) + args.append(tvm.tir.Var(name, default_index_dtype)) + else: + raise ValueError("transform_layout mapping may not have *args or **kwargs") + final_indices = mapping_function(*args) + return IndexMap(args, final_indices) diff --git a/src/tir/ir/index_map.cc b/src/tir/ir/index_map.cc index 7955c4f7fc6f..d8b1123fdd9d 100644 --- a/src/tir/ir/index_map.cc +++ b/src/tir/ir/index_map.cc @@ -189,7 +189,7 @@ String IndexMapNode::ToPythonString() const { TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) .set_dispatch([](const ObjectRef& node, ReprPrinter* p) { auto* op = static_cast(node.get()); - p->stream << "index_map(" << op->initial_indices << ", " << op->final_indices << ")"; + p->stream << "index_map(" << op->ToPythonString() << ")"; }); TVM_REGISTER_NODE_TYPE(IndexMapNode); diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index 8d489fe601b4..914f0426ec49 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -27,15 +27,14 @@ # fmt: off # pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks + def packed_index_map_func(m, n): return m // 16, n // 16, m % 16, n % 16 @T.prim_func -def two_elementwise(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128), "float32") +def two_elementwise(A: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"]) -> None: B = T.alloc_buffer((128, 128), "float32") - C = T.match_buffer(c, (128, 128), "float32") for i, j in T.grid(128, 128): with T.block("B"): vi, vj = T.axis.remap("SS", [i, j]) @@ -47,10 +46,10 @@ def two_elementwise(a: T.handle, c: T.handle) -> None: @T.prim_func -def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128), "float32") +def two_elementwise_transformed_intermediate_buffer( + A: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"] +) -> None: B = T.alloc_buffer((8, 8, 16, 16), "float32") - C = T.match_buffer(c, (128, 128), "float32") for i, j in T.grid(128, 128): with T.block("B"): vi, vj = T.axis.remap("SS", [i, j]) @@ -62,10 +61,10 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> @T.prim_func -def two_elementwise_transformed_input_buffer(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (8, 8, 16, 16), "float32") +def two_elementwise_transformed_input_buffer( + A: T.Buffer[(8, 8, 16, 16), "float32"], C: T.Buffer[(128, 128), "float32"] +) -> None: B = T.alloc_buffer((128, 128), "float32") - C = T.match_buffer(c, (128, 128), "float32") for i, j in T.grid(128, 128): with T.block("B"): vi, vj = T.axis.remap("SS", [i, j]) @@ -77,10 +76,10 @@ def two_elementwise_transformed_input_buffer(a: T.handle, c: T.handle) -> None: @T.prim_func -def two_elementwise_transformed_output_buffer(a: T.handle, c: T.handle) -> None: - A = T.match_buffer(a, (128, 128), "float32") +def two_elementwise_transformed_output_buffer( + A: T.Buffer[(128, 128), "float32"], C: T.Buffer[(8, 8, 16, 16), "float32"] +) -> None: B = T.alloc_buffer((128, 128), "float32") - C = T.match_buffer(c, (8, 8, 16, 16), "float32") for i, j in T.grid(128, 128): with T.block("B"): vi, vj = T.axis.remap("SS", [i, j]) From 022ddd3a1264d55ae647444822c6b87ed6029d2c Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 14 Mar 2022 19:52:04 -0400 Subject: [PATCH 06/11] remove unused --- src/tir/ir/index_map.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/tir/ir/index_map.cc b/src/tir/ir/index_map.cc index d8b1123fdd9d..6185f1097a76 100644 --- a/src/tir/ir/index_map.cc +++ b/src/tir/ir/index_map.cc @@ -199,7 +199,5 @@ TVM_REGISTER_GLOBAL("tir.IndexMap") return IndexMap(initial_indices, final_indices); }); -TVM_REGISTER_GLOBAL("tir.IndexMapFromFunc").set_body_typed(IndexMap::FromFunc); - } // namespace tir } // namespace tvm From 21bcab5dd1b6aae1165caa889df3fd16e253b982 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 21 Mar 2022 19:59:21 -0400 Subject: [PATCH 07/11] Use BufferIndexType enum --- include/tvm/tir/schedule/schedule.h | 14 +++++++++++--- python/tvm/tir/__init__.py | 2 +- python/tvm/tir/schedule/__init__.py | 2 +- python/tvm/tir/schedule/schedule.py | 18 +++++++++++++----- src/tir/schedule/concrete_schedule.cc | 5 +++-- src/tir/schedule/concrete_schedule.h | 2 +- src/tir/schedule/primitive.h | 4 ++-- .../primitive/layout_transformation.cc | 13 ++++++------- src/tir/schedule/schedule.cc | 6 +++++- src/tir/schedule/traced_schedule.cc | 7 ++++--- src/tir/schedule/traced_schedule.h | 2 +- .../test_tir_schedule_transform_layout.py | 7 ++++--- 12 files changed, 52 insertions(+), 30 deletions(-) diff --git a/include/tvm/tir/schedule/schedule.h b/include/tvm/tir/schedule/schedule.h index 1507955f5cfc..0273ece0b3b1 100644 --- a/include/tvm/tir/schedule/schedule.h +++ b/include/tvm/tir/schedule/schedule.h @@ -37,6 +37,14 @@ enum class ScheduleErrorRenderLevel : int32_t { kNone = 2, }; +/*! \brief Type of buffer index */ +enum class BufferIndexType : int32_t { + /*! \brief Index of a read buffer */ + kRead = 0, + /*! \brief Index of a written buffer */ + kWrite = 1, +}; + /**************** Random variable: BlockRV ****************/ /*! \brief A random variable that evaluates to a TensorIR block */ @@ -531,11 +539,11 @@ class ScheduleNode : public runtime::Object { * 'match_buffer'). * \param block_rv The block that accesses the target buffer. * \param buffer_index The index of the buffer in block's read or write region. - * \param is_write_index Whether the buffer_index is the index of the block's write region. + * \param buffer_index_type The type of the buffer index, kRead or kWrite. * \param index_map The transformation to apply. */ - virtual void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, - const IndexMap& index_map) = 0; + virtual void TransformLayout(const BlockRV& block_rv, int buffer_index, + BufferIndexType buffer_index_type, const IndexMap& index_map) = 0; /******** Schedule: Misc ********/ /*! \brief A no-op that marks the start of postprocessing phase of scheduling */ diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py index 17f9aa3d9c60..147360d7e087 100644 --- a/python/tvm/tir/__init__.py +++ b/python/tvm/tir/__init__.py @@ -57,7 +57,7 @@ from .op import comm_reducer, min, max, sum from .op import q_multiply_shift -from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError +from .schedule import StmtSRef, BlockScope, ScheduleState, Schedule, ScheduleError, BufferType from . import schedule from . import ir_builder diff --git a/python/tvm/tir/schedule/__init__.py b/python/tvm/tir/schedule/__init__.py index 5f0e169c43e3..2314c7fb939f 100644 --- a/python/tvm/tir/schedule/__init__.py +++ b/python/tvm/tir/schedule/__init__.py @@ -19,6 +19,6 @@ from .block_scope import BlockScope, Dependency, DepKind, StmtSRef from .instruction import Instruction, InstructionKind -from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError +from .schedule import BlockRV, ExprRV, LoopRV, Schedule, ScheduleError, BufferType from .state import ScheduleDebugMask, ScheduleState from .trace import Trace diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 622ca6c4c40d..8532c6e8daf5 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -28,6 +28,7 @@ from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod from .trace import Trace from ._type_checker import type_checked +import enum @register_error @@ -72,6 +73,13 @@ def __init__(self) -> None: } +class BufferType(enum.IntEnum): + """Type of buffer in access regions of a block""" + + READ = 0 + WRITE = 1 + + def _parse_error_render_level(error_render_level: str) -> int: if error_render_level not in _ERROR_RENDER_LEVEL: raise ValueError( @@ -2119,7 +2127,7 @@ def transform_layout( self, block: BlockRV, buffer_index: int, - is_write_index: bool, + buffer_type: BufferType, index_map: Union[IndexMap, Callable], ) -> None: """Apply a transformation represented by IndexMap to buffer @@ -2129,8 +2137,8 @@ def transform_layout( The block that accesses the target buffer buffer_index: int The index of the buffer in block's read or write region - is_write_index : bool - Whether the buffer_index is the index of the block's write region + buffer_type : BufferType + Type of the buffer, READ or WRITE. index_map : Union[IndexMap, Callable] The transformation to apply @@ -2159,7 +2167,7 @@ def before_transform_layout(a: T.handle, c: T.handle) -> None: .. code-block:: python sch = tir.Schedule(before_storage_align) - sch.transform_layout(sch.get_block("B"), buffer_index=0, is_write_index=True, + sch.transform_layout(sch.get_block("B"), buffer_index=0, BufferType.WRITE, index_map=lambda m, n: (m // 16, n // 16, m % 16, n % 16)) print(sch.mod["main"].script()) @@ -2185,7 +2193,7 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> if callable(index_map): index_map = IndexMap.from_func(index_map) _ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member - self, block, buffer_index, is_write_index, index_map + self, block, buffer_index, buffer_type, index_map ) ########## Schedule: Misc ########## diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index e1cf46bfe6ec..331ae0209cc0 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -687,9 +687,10 @@ void ConcreteScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann /******** Schedule: Layout transformation ********/ void ConcreteScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, - bool is_write_index, const IndexMap& index_map) { + BufferIndexType buffer_index_type, + const IndexMap& index_map) { TVM_TIR_SCHEDULE_BEGIN(); - tir::TransformLayout(state_, this->GetSRef(block_rv), buffer_index, is_write_index, index_map); + tir::TransformLayout(state_, this->GetSRef(block_rv), buffer_index, buffer_index_type, index_map); 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 f695d523b46c..32aab1a7b44d 100644 --- a/src/tir/schedule/concrete_schedule.h +++ b/src/tir/schedule/concrete_schedule.h @@ -132,7 +132,7 @@ class ConcreteScheduleNode : public ScheduleNode { void Annotate(const BlockRV& block_rv, const String& ann_key, const ObjectRef& ann_val) override; void Unannotate(const BlockRV& block_rv, const String& ann_key) override; /******** Schedule: Layout transformation ********/ - void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, + void TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map) override; /******** Schedule: Misc ********/ void EnterPostproc() override {} diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h index b32b99fb446b..5e21075d5844 100644 --- a/src/tir/schedule/primitive.h +++ b/src/tir/schedule/primitive.h @@ -425,11 +425,11 @@ TVM_DLL void Unannotate(ScheduleState self, const StmtSRef& sref, const String& * \param self The state of the schedule * \param block_sref The block sref that accesses the target buffer. * \param buffer_index The index of the buffer in block's read or write region. - * \param is_write_index Whether the buffer_index is the index of the block's write region. + * \param buffer_index_type The type of the buffer index, kRead or kWrite. * \param index_map The transformation to apply. */ TVM_DLL void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, - bool is_write_index, const IndexMap& index_map); + BufferIndexType buffer_index_type, const IndexMap& index_map); /******** Schedule: Misc ********/ diff --git a/src/tir/schedule/primitive/layout_transformation.cc b/src/tir/schedule/primitive/layout_transformation.cc index 4570ab43ee9a..e7a57b18ca73 100644 --- a/src/tir/schedule/primitive/layout_transformation.cc +++ b/src/tir/schedule/primitive/layout_transformation.cc @@ -126,10 +126,9 @@ class BufferIsSubregionError : public ScheduleError { }; void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, - bool is_write_index, const IndexMap& index_map) { + BufferIndexType buffer_index_type, const IndexMap& index_map) { const BlockNode* block_ptr = TVM_SREF_TO_BLOCK(block_ptr, block_sref); - Buffer old_buffer = GetNthAccessBuffer(self, GetRef(block_ptr), buffer_index, - /*is_write=*/is_write_index); + Buffer old_buffer = GetNthAccessBuffer(self, GetRef(block_ptr), buffer_index, buffer_index_type == BufferIndexType::kRead ? false : true); Optional defining_site_sref; bool is_alloc; std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, old_buffer); @@ -196,16 +195,16 @@ struct TransformLayoutTraits : public UnpackedInstTraits static constexpr size_t kNumDecisions = 0; static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv, Integer buffer_index, - Bool is_write_index, IndexMap index_map) { - return sch->TransformLayout(block_rv, buffer_index, is_write_index, index_map); + Integer buffer_index_type, IndexMap index_map) { + return sch->TransformLayout(block_rv, buffer_index, static_cast(buffer_index_type->value), index_map); } static String UnpackedAsPython(Array outputs, String block_rv, Integer buffer_index, - Bool is_write_index, IndexMap index_map) { + Integer buffer_index_type, IndexMap index_map) { PythonAPICall py("transform_layout"); py.Input("block", block_rv); py.Input("buffer_index", buffer_index); - py.Input("is_write_index", is_write_index); + py.Input("buffer_index_type", buffer_index_type); py.Input("index_map", index_map->ToPythonString()); return py.Str(); } diff --git a/src/tir/schedule/schedule.cc b/src/tir/schedule/schedule.cc index 060f9fbe7619..82cd0a4a351a 100644 --- a/src/tir/schedule/schedule.cc +++ b/src/tir/schedule/schedule.cc @@ -228,7 +228,11 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnannotate") /******** (FFI) Layout transformation ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformLayout") - .set_body_method(&ScheduleNode::TransformLayout); + .set_body_typed([](Schedule self, const BlockRV& block_rv, int buffer_index, + int buffer_index_type, const IndexMap& index_map) { + return self->TransformLayout(block_rv, buffer_index, + static_cast(buffer_index_type), index_map); + }); /******** (FFI) Misc ********/ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleEnterPostproc") diff --git a/src/tir/schedule/traced_schedule.cc b/src/tir/schedule/traced_schedule.cc index 1d41d9e443eb..8af66f1ede75 100644 --- a/src/tir/schedule/traced_schedule.cc +++ b/src/tir/schedule/traced_schedule.cc @@ -430,13 +430,14 @@ void TracedScheduleNode::Unannotate(const BlockRV& block_rv, const String& ann_k /******** Schedule: Layout transformation ********/ void TracedScheduleNode::TransformLayout(const BlockRV& block_rv, int buffer_index, - bool is_write_index, const IndexMap& index_map) { - ConcreteScheduleNode::TransformLayout(block_rv, buffer_index, is_write_index, index_map); + BufferIndexType buffer_index_type, + const IndexMap& index_map) { + ConcreteScheduleNode::TransformLayout(block_rv, buffer_index, buffer_index_type, index_map); static const InstructionKind& kind = InstructionKind::Get("TransformLayout"); trace_->Append( /*inst=*/Instruction(/*kind=*/kind, /*inputs=*/{block_rv}, - /*attrs=*/{Integer(buffer_index), Bool(is_write_index), index_map}, + /*attrs=*/{Integer(buffer_index), Integer(buffer_index_type), index_map}, /*outputs=*/{})); } diff --git a/src/tir/schedule/traced_schedule.h b/src/tir/schedule/traced_schedule.h index 9ad9b237bc1a..5d355bd70c99 100644 --- a/src/tir/schedule/traced_schedule.h +++ b/src/tir/schedule/traced_schedule.h @@ -96,7 +96,7 @@ class TracedScheduleNode : public ConcreteScheduleNode { void Annotate(const BlockRV& block_rv, const String& ann_key, const ObjectRef& ann_val) override; void Unannotate(const BlockRV& block_rv, const String& ann_key) override; /******** Schedule: Layout transformation ********/ - void TransformLayout(const BlockRV& block_rv, int buffer_index, bool is_write_index, + void TransformLayout(const BlockRV& block_rv, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map) override; /******** Schedule: Misc ********/ void EnterPostproc() final; diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index 914f0426ec49..1c240eb7d748 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -21,6 +21,7 @@ import tvm from tvm import tir +from tvm.tir import BufferType from tvm.script import tir as T from tvm.tir.schedule.testing import verify_trace_roundtrip @@ -97,7 +98,7 @@ def two_elementwise_transformed_output_buffer( def test_two_elementwise_transform_intermediate_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, True, lambda m, n: (m // 16, n // 16, m % 16, n % 16)) + sch.transform_layout(block, 0, BufferType.WRITE, lambda m, n: (m // 16, n // 16, m % 16, n % 16)) tvm.ir.assert_structural_equal(two_elementwise_transformed_intermediate_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) @@ -105,7 +106,7 @@ def test_two_elementwise_transform_intermediate_buffer(): def test_two_elementwise_transform_input_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, False, packed_index_map_func) + sch.transform_layout(block, 0, BufferType.READ, packed_index_map_func) tvm.ir.assert_structural_equal(two_elementwise_transformed_input_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) @@ -113,7 +114,7 @@ def test_two_elementwise_transform_input_buffer(): def test_two_elementwise_transform_output_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("C") - sch.transform_layout(block, 0, True, packed_index_map_func) + sch.transform_layout(block, 0, BufferType.WRITE, packed_index_map_func) tvm.ir.assert_structural_equal(two_elementwise_transformed_output_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) From dcdf22c574384093e927508088cf06779afef5b1 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 21 Mar 2022 20:32:51 -0400 Subject: [PATCH 08/11] lint --- src/tir/schedule/primitive/layout_transformation.cc | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/tir/schedule/primitive/layout_transformation.cc b/src/tir/schedule/primitive/layout_transformation.cc index e7a57b18ca73..56eedca1120d 100644 --- a/src/tir/schedule/primitive/layout_transformation.cc +++ b/src/tir/schedule/primitive/layout_transformation.cc @@ -128,7 +128,9 @@ class BufferIsSubregionError : public ScheduleError { void TransformLayout(ScheduleState self, const StmtSRef& block_sref, int buffer_index, BufferIndexType buffer_index_type, const IndexMap& index_map) { const BlockNode* block_ptr = TVM_SREF_TO_BLOCK(block_ptr, block_sref); - Buffer old_buffer = GetNthAccessBuffer(self, GetRef(block_ptr), buffer_index, buffer_index_type == BufferIndexType::kRead ? false : true); + Buffer old_buffer = + GetNthAccessBuffer(self, GetRef(block_ptr), buffer_index, + buffer_index_type == BufferIndexType::kRead ? false : true); Optional defining_site_sref; bool is_alloc; std::tie(defining_site_sref, is_alloc) = GetBufferDefiningSite(block_sref, old_buffer); @@ -196,7 +198,8 @@ struct TransformLayoutTraits : public UnpackedInstTraits static void UnpackedApplyToSchedule(Schedule sch, BlockRV block_rv, Integer buffer_index, Integer buffer_index_type, IndexMap index_map) { - return sch->TransformLayout(block_rv, buffer_index, static_cast(buffer_index_type->value), index_map); + return sch->TransformLayout(block_rv, buffer_index, + static_cast(buffer_index_type->value), index_map); } static String UnpackedAsPython(Array outputs, String block_rv, Integer buffer_index, From a83c41393c1db393a9615b481f3888ccb3f12df9 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Mon, 21 Mar 2022 20:47:52 -0400 Subject: [PATCH 09/11] support *args --- python/tvm/tir/function.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py index 031ab7ce2c68..98af3b472030 100644 --- a/python/tvm/tir/function.py +++ b/python/tvm/tir/function.py @@ -16,7 +16,7 @@ # under the License. """Function data types.""" -from typing import Callable, List, Mapping, Union +from typing import Callable, List, Mapping, Optional, Union import inspect import tvm._ffi @@ -261,7 +261,7 @@ def __init__(self, initial_indices, final_indices): self.__init_handle_by_constructor__(_ffi_api.IndexMap, initial_indices, final_indices) @staticmethod - def from_func(mapping_function: Callable): + def from_func(mapping_function: Callable, ndim: Optional[int] = None): """Create an index map from a function Parameters @@ -272,14 +272,26 @@ def from_func(mapping_function: Callable): params = inspect.signature(mapping_function).parameters default_index_dtype = "int32" args = [] + var_arg_name = None for name, param in params.items(): if param.kind in [ inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, ]: - print(name) args.append(tvm.tir.Var(name, default_index_dtype)) + elif param.kind == inspect.Parameter.VAR_POSITIONAL: + var_arg_name = name else: raise ValueError("transform_layout mapping may not have *args or **kwargs") + + # Now that all the named arguments have been collected, + # everything that remains should go to the *args, if + # specified. + if var_arg_name is not None: + assert ndim is not None, "ndim must be specified when *args is used" + num_var_args = ndim - len(args) + for i in range(num_var_args): + args.append(tvm.tir.Var(f"{var_arg_name}_{i}", default_index_dtype)) + final_indices = mapping_function(*args) return IndexMap(args, final_indices) From d518267d793e0ab322d601edf5e05cd1355d3367 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Tue, 22 Mar 2022 14:25:46 -0400 Subject: [PATCH 10/11] lint --- tests/python/unittest/test_tir_schedule_transform_layout.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/unittest/test_tir_schedule_transform_layout.py b/tests/python/unittest/test_tir_schedule_transform_layout.py index 1c240eb7d748..e0a7f66bf278 100644 --- a/tests/python/unittest/test_tir_schedule_transform_layout.py +++ b/tests/python/unittest/test_tir_schedule_transform_layout.py @@ -98,7 +98,9 @@ def two_elementwise_transformed_output_buffer( def test_two_elementwise_transform_intermediate_buffer(): sch = tir.Schedule(two_elementwise, debug_mask="all") block = sch.get_block("B") - sch.transform_layout(block, 0, BufferType.WRITE, lambda m, n: (m // 16, n // 16, m % 16, n % 16)) + sch.transform_layout( + block, 0, BufferType.WRITE, lambda m, n: (m // 16, n // 16, m % 16, n % 16) + ) tvm.ir.assert_structural_equal(two_elementwise_transformed_intermediate_buffer, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=two_elementwise) From d82bccf26a167bf95acbeb6b453c62cf4536a9eb Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Tue, 22 Mar 2022 16:11:36 -0400 Subject: [PATCH 11/11] lint --- python/tvm/tir/schedule/schedule.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py index 8532c6e8daf5..c54c7f74f24f 100644 --- a/python/tvm/tir/schedule/schedule.py +++ b/python/tvm/tir/schedule/schedule.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. """The TensorIR schedule class""" +import enum from typing import Callable, Dict, List, Optional, Union from tvm._ffi import register_object as _register_object @@ -28,7 +29,6 @@ from .state import ScheduleState, StmtSRef, _parse_debug_mask, _parse_mod from .trace import Trace from ._type_checker import type_checked -import enum @register_error