wehub-resource-sync
939 lines
45 KiB
C++
939 lines
45 KiB
C++
/*
|
|
* 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.
|
|
*/
|
|
#ifndef TVM_S_TIR_SCHEDULE_SCHEDULE_H_
|
|
#define TVM_S_TIR_SCHEDULE_SCHEDULE_H_
|
|
|
|
#include <tvm/s_tir/random_engine.h>
|
|
#include <tvm/s_tir/schedule/state.h>
|
|
#include <tvm/s_tir/schedule/trace.h>
|
|
#include <tvm/tirx/index_map.h>
|
|
|
|
namespace tvm {
|
|
namespace s_tir {
|
|
using namespace tvm::tirx;
|
|
|
|
/*! \brief The level of detailed error message rendering */
|
|
enum class ScheduleErrorRenderLevel : int32_t {
|
|
/*! \brief Render a detailed error message */
|
|
kDetail = 0,
|
|
/*! \brief Render the error in fast mode */
|
|
kFast = 1,
|
|
/*! \brief No error message at all */
|
|
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: SBlockRV ****************/
|
|
|
|
/*! \brief A random variable that evaluates to a TensorIR block */
|
|
class SBlockRVNode : public ffi::Object {
|
|
public:
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<SBlockRVNode>();
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBlockRV", SBlockRVNode, ffi::Object);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to SBlockRVNode
|
|
* \sa SBlockRVNode
|
|
*/
|
|
class SBlockRV : public ffi::ObjectRef {
|
|
public:
|
|
/*! \brief Constructor */
|
|
TVM_DLL SBlockRV();
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SBlockRV, ffi::ObjectRef, SBlockRVNode);
|
|
};
|
|
|
|
/**************** Random variable: LoopRV ****************/
|
|
|
|
/*! \brief A random variable that evaluates to a TensorIR for loop */
|
|
class LoopRVNode : public ffi::Object {
|
|
public:
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<LoopRVNode>();
|
|
}
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.LoopRV", LoopRVNode, ffi::Object);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to LoopRVNode
|
|
* \sa LoopRVNode
|
|
*/
|
|
class LoopRV : public ffi::ObjectRef {
|
|
public:
|
|
/*! \brief Constructor */
|
|
TVM_DLL LoopRV();
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LoopRV, ffi::ObjectRef, LoopRVNode);
|
|
};
|
|
|
|
/**************** Random variable: ExprRV ****************/
|
|
|
|
/*! \brief An expr random variable */
|
|
using ExprRV = PrimExpr;
|
|
|
|
using ExprRVNode = ExprNode;
|
|
|
|
/**************** The Schedule class ****************/
|
|
|
|
class Schedule;
|
|
|
|
/*! \brief The user-facing schedule class */
|
|
class ScheduleNode : public ffi::Object {
|
|
friend class Schedule;
|
|
|
|
public:
|
|
virtual ~ScheduleNode() = default;
|
|
|
|
static constexpr const bool _type_mutable = true;
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Schedule", ScheduleNode, ffi::Object);
|
|
|
|
public:
|
|
/*! \brief Get the IRModule associated with this schedule. */
|
|
virtual IRModule mod() const { return state()->mod; }
|
|
/*! \return The internal state of scheduling */
|
|
virtual ScheduleState state() const = 0;
|
|
/*! \return The internally maintained trace of scheduling program execution */
|
|
virtual ffi::Optional<Trace> trace() const = 0;
|
|
/*! \return The GlobalVar of the func that the schedule is currently working on */
|
|
virtual ffi::Optional<GlobalVar> func_working_on() const = 0;
|
|
/*!
|
|
* \brief Instruct the schedule to work on a function in the IRModule.
|
|
*
|
|
* By default, the schedule works on the function with the name "main", or the only function in
|
|
* the IRModule if there is only one. If there is multiple functions in the IRModule, and none of
|
|
* their names are "main", users will have to call this method to explicitly specify which
|
|
* function to work on.
|
|
*
|
|
* This sugar function will guide the `GetSBlock` method if its `func_name` is not specified.
|
|
*
|
|
* \param func_name The name of the function to be working on
|
|
*
|
|
* \sa GetSBlock
|
|
*/
|
|
virtual void WorkOn(const ffi::String& func_name) = 0;
|
|
/*!
|
|
* \brief Returns a copy of the schedule, including both its state and its symbol table,
|
|
* guaranteeing that
|
|
* 1) SRef tree is completely reconstructed;
|
|
* 2) The IRModule being scheduled is not modified;
|
|
* 3) All the random variables are valid in the copy, pointing to the corresponding sref
|
|
* reconstructed
|
|
*/
|
|
virtual Schedule Copy() = 0;
|
|
/*!
|
|
* \brief Seed the randomness
|
|
* \param seed The new random seed, -1 if use device random, otherwise non-negative
|
|
*/
|
|
virtual void Seed(LinearCongruentialEngine::TRandState seed) = 0;
|
|
/*! \brief Fork the random state */
|
|
virtual LinearCongruentialEngine::TRandState ForkSeed() = 0;
|
|
|
|
public:
|
|
/******** Lookup/Remove random variables ********/
|
|
/*!
|
|
* \brief Get the block corresponding to the specific SBlockRV
|
|
* \param block_rv The SBlockRV to be looked up
|
|
* \return The corresponding block
|
|
*/
|
|
virtual SBlock Get(const SBlockRV& block_rv) const = 0;
|
|
/*!
|
|
* \brief Get the for loop corresponding to the specific LoopRV
|
|
* \param loop_rv The LoopRV to be looked up
|
|
* \return The corresponding for loop
|
|
*/
|
|
virtual For Get(const LoopRV& loop_rv) const = 0;
|
|
/*!
|
|
* \brief Get the expr corresponding to the specific random variable
|
|
* \param expr_rv The random variable to be looked up
|
|
* \return The corresponding expr
|
|
*/
|
|
virtual PrimExpr Get(const ExprRV& expr_rv) const = 0;
|
|
/*!
|
|
* \brief Get the block sref corresponding to the specific SBlockRV
|
|
* \param block_rv The SBlockRV to be looked up
|
|
* \return The corresponding block sref
|
|
*/
|
|
virtual StmtSRef GetSRef(const SBlockRV& block_rv) const = 0;
|
|
/*!
|
|
* \brief Get the loop sref corresponding to the specific LoopRV
|
|
* \param loop_rv The LoopRV to be looked up
|
|
* \return The corresponding loop sref
|
|
*/
|
|
virtual StmtSRef GetSRef(const LoopRV& loop_rv) const = 0;
|
|
/*!
|
|
* \brief Check the existance of a specific SBlockRV
|
|
* \param block_rv The SBlockRV to be looked up
|
|
* \return Whether the corresponding block exists
|
|
*/
|
|
virtual bool HasBlock(const SBlockRV& block_rv) const = 0;
|
|
/*!
|
|
* \brief Get the block/loop sref corresponding to the specific statement
|
|
* \param stmt The statement to be looked up
|
|
* \return The corresponding block/loop sref
|
|
*/
|
|
virtual StmtSRef GetSRef(const StmtNode* stmt) const;
|
|
/*!
|
|
* \brief Get the block/loop sref corresponding to the specific statement
|
|
* \param stmt The statement to be looked up
|
|
* \return The corresponding block/loop sref
|
|
*/
|
|
StmtSRef GetSRef(const Stmt& stmt) const { return this->GetSRef(stmt.get()); }
|
|
/*!
|
|
* \brief Remove a block random variable from the symbol table
|
|
* \param block_rv The random variable to be removed
|
|
*/
|
|
virtual void RemoveRV(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Remove a loop random variable from the symbol table
|
|
* \param loop_rv The random variable to be removed
|
|
*/
|
|
virtual void RemoveRV(const LoopRV& loop_rv) = 0;
|
|
/*!
|
|
* \brief Remove an integer random variable from the symbol table
|
|
* \param expr_rv The random variable to be removed
|
|
*/
|
|
virtual void RemoveRV(const ExprRV& expr_rv) = 0;
|
|
|
|
public:
|
|
/******** Schedule: Sampling ********/
|
|
/*!
|
|
* \brief Sample an integer given the probability distribution
|
|
* \param candidates The candidates
|
|
* \param probs The probability distribution of the candidates
|
|
* \param decision The sampling decision
|
|
* \return The random variable sampled from candidates
|
|
*/
|
|
virtual ExprRV SampleCategorical(const ffi::Array<int64_t>& candidates,
|
|
const ffi::Array<FloatImm>& probs,
|
|
ffi::Optional<int64_t> decision = std::nullopt) = 0;
|
|
/*!
|
|
* \brief Sample the factors to perfect tile a specific loop
|
|
* \param loop_rv The loop to be tiled
|
|
* \param n The number of tiles to be sampled
|
|
* \param max_innermost_factor The maximum tile size allowed to be sampled in the innermost loop
|
|
* \param decision The sampling decision
|
|
* \return A list of length `n`, the random perfect tile sizes sampled
|
|
*/
|
|
virtual ffi::Array<ExprRV> SamplePerfectTile(
|
|
const LoopRV& loop_rv, int n, int max_innermost_factor,
|
|
ffi::Optional<ffi::Array<int64_t>> decision = std::nullopt) = 0;
|
|
/*!
|
|
* \brief Sample the factors to a partitioned tile for a specific loop
|
|
*
|
|
* The sampled tile size will be partitioned into two parts. The second part has a guarantee
|
|
* that their extent's product have a factor of `innerpart_factor`. The first part is loops at
|
|
* [0, partition_pos); the second part is loops at [partition_pos, n) and we will have
|
|
* `innerpart_factor` | (l[partition_pos].extent * ... * l[n-1].extent)
|
|
*
|
|
* \param loop_rv The loop to be tiled
|
|
* \param n The number of tiles to be sampled
|
|
* \param partition_pos The position to partition tiles to two parts
|
|
* \param innerpart_factor The factor of the second part
|
|
* \param decision The sampling decision
|
|
* \return A list of length `n`, the random partitioned tile sizes sampled
|
|
*/
|
|
virtual ffi::Array<ExprRV> SamplePartitionedTile(
|
|
const LoopRV& loop_rv, int n, int partition_pos, int innerpart_factor,
|
|
ffi::Optional<ffi::Array<int64_t>> decision = std::nullopt) = 0;
|
|
/*!
|
|
* \brief Sample a compute-at location of the given block
|
|
* \param block_rv The block whose compute-at location is to be sampled
|
|
* \param decision The sampling decision
|
|
* \return The sampled loop where the input block is to be computed at
|
|
*/
|
|
virtual LoopRV SampleComputeLocation(const SBlockRV& block_rv,
|
|
ffi::Optional<int64_t> decision = std::nullopt) = 0;
|
|
|
|
/******** Schedule: Get blocks & loops ********/
|
|
/*!
|
|
* \brief Retrieve a block in a specific function with its name
|
|
*
|
|
* By default, if `func_name` is not specified, the schedule will search for the block in the
|
|
* function that is currently being "worked on". To switch the function to be worked on, use
|
|
* `WorkOn` before calling this method.
|
|
*
|
|
* \param name The name of the block to be retrieved
|
|
* \param func_name The name of the function
|
|
* \return The block retrieved
|
|
* \note Indexing error is raised if 0 or multiple blocks exist with the specific name
|
|
*
|
|
* \sa WorkOn
|
|
*/
|
|
virtual SBlockRV GetSBlock(const ffi::String& name,
|
|
const ffi::Optional<ffi::String>& func_name = std::nullopt) = 0;
|
|
/*!
|
|
* \brief Get the parent loops of the block in its scope, from outer to inner
|
|
* \param block_rv The query block
|
|
* \return A list of loops above the given block in its scope, from outer to inner
|
|
*/
|
|
virtual ffi::Array<LoopRV> GetLoops(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Get the leaf blocks of a specific scope
|
|
* \param block_rv The block where the scope is rooted
|
|
* \return A list of child blocks
|
|
*/
|
|
virtual ffi::Array<SBlockRV> GetChildBlocks(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Get the leaf blocks of under a specific loop
|
|
* \param loop_rv The loop under which collecting is conducted
|
|
* \return A list of child blocks
|
|
*/
|
|
virtual ffi::Array<SBlockRV> GetChildBlocks(const LoopRV& loop_rv) = 0;
|
|
/*!
|
|
* \brief Get the producer of a specific block, under the same block scope
|
|
* \param block_rv The block in the query
|
|
* \return A list of blocks, the producers of the given block under the same scope of the given
|
|
* block
|
|
*/
|
|
virtual ffi::Array<SBlockRV> GetProducers(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Get the consumers of a specific block, under the same block scope
|
|
* \param block_rv The block to be queried
|
|
* \return A list of blocks, the consumers of the given block under the same scope of the given
|
|
* block
|
|
*/
|
|
virtual ffi::Array<SBlockRV> GetConsumers(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Get the list of output blocks within the given scope
|
|
* An output block is a block which has atleast one buffer being written
|
|
* to, but is not allocated within the PrimFunc
|
|
* \param scope_block_rv The scope block from which output blocks are collected
|
|
* \return A list of all blocks that write to some output buffer
|
|
* block
|
|
*/
|
|
virtual ffi::Array<SBlockRV> GetOutputBlocks(const SBlockRV& scope_block_rv) = 0;
|
|
/******** Schedule: Transform loops ********/
|
|
/*!
|
|
* \brief Merge a list of loops into one. The loops under their LCA requires:
|
|
* 1) Under the same scope
|
|
* 2) Can't have annotations or thread bindings
|
|
* 3) Start with 0 and have same extent and same nesting depth
|
|
* 4) From target loop to their LCA, the inner loop must be the only child of the outer loop
|
|
* \param loop_rvs The loops to be merged
|
|
* \return The new loop after merge
|
|
*/
|
|
virtual LoopRV Merge(const ffi::Array<LoopRV>& loop_rvs) = 0;
|
|
/*!
|
|
* \brief Fuse a list of consecutive loops into one. It requires:
|
|
* 1) The loops can't have annotations or thread bindings.
|
|
* 2) The (i+1)-th loop must be the only child of the i-th loop.
|
|
* 3) All loops must start with 0.
|
|
* 4) The domain of a loop to be fused cannot depend on another loop to be fused.
|
|
* \param loop_rvs The loops to be fused
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
* \return The new loop after fusion
|
|
*/
|
|
virtual LoopRV Fuse(const ffi::Array<LoopRV>& loop_rvs, bool preserve_unit_iters = true) = 0;
|
|
/*!
|
|
* \brief Split a loop into a list of consecutive loops. It requires:
|
|
* 1) The loop can't have annotation or thread binding.
|
|
* 2) The loop must start with 0.
|
|
* \param loop_rv The loop to be split
|
|
* \param factors The positive tiling factors, and at most one of which is `std::nullopt`, which
|
|
* means that factor is inferred. \param preserve_unit_iters Whether or not to preserve unit
|
|
* iterators in block bindings. \param disable_predication If enabled, don't create a predicate
|
|
* for guarding the loop. This can be useful when splitting with scalable factors that the
|
|
* schedule writer knows are divisible by the loop bound. Warning: enabling this feature may
|
|
* result in incorrect code generation if not used carefully. \return The new loops after split.
|
|
*/
|
|
virtual ffi::Array<LoopRV> Split(const LoopRV& loop_rv,
|
|
const ffi::Array<ffi::Optional<ExprRV>>& factors,
|
|
bool preserve_unit_iters = true,
|
|
bool disable_predication = false) = 0;
|
|
/*!
|
|
* \brief Partition the loops into sequence of multiple loops
|
|
* 1) The loop can't have annotation or thread binding.
|
|
* \param loop_rv The loop to be partition
|
|
* \param factors The positive integers, and at most one of which is `std::nullopt`, which means
|
|
* that factor is inferred.
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
* \return The new loops after partition
|
|
*/
|
|
virtual ffi::Array<LoopRV> LoopPartition(const LoopRV& loop_rv,
|
|
const ffi::Array<ffi::Optional<ExprRV>>& factors,
|
|
bool preserve_unit_iters = true) = 0;
|
|
/*!
|
|
* \brief Reorder a list of loops. It doesn't require the loops to be consecutive.
|
|
* It requires:
|
|
* 1) The loops are in the same chain. That means: the loops can be ordered to [l_1, l_2, ... ,
|
|
* l_n] where l_i is an ancestor of l_{i+1} and there are only single-branch loops between
|
|
* l_1 and l_n (which also indicates they are under the same scope).
|
|
* 2) After reordering, the domain of an outer loop cannot depend on any of the inner loops.
|
|
* 3) For every block under the loop nests, its block binding must be affine, and the block
|
|
* variables must be either data parallel or reduction.
|
|
* 4) No duplicated loops are allowed in the arguments.
|
|
* \param ordered_loop_rvs The loops in the new order
|
|
*/
|
|
virtual void Reorder(const ffi::Array<LoopRV>& ordered_loop_rvs) = 0;
|
|
/*!
|
|
* \brief Reorder the itervars inside a block.
|
|
* \param block_rv The block to be transformed.
|
|
* \param new_order The new itervar order.
|
|
*/
|
|
virtual void ReorderBlockIterVar(const SBlockRV& block_rv,
|
|
const ffi::Array<int64_t> new_order) = 0;
|
|
/*!
|
|
* \brief Create a new unit loop on top of the specific block.
|
|
* \param block_rv The block above which the new loop is created
|
|
* \return The new loop created
|
|
*/
|
|
virtual LoopRV AddUnitLoop(const SBlockRV& block_rv) = 0;
|
|
/*!
|
|
* \brief Create a new unit loop on top of the specific loop.
|
|
* \param loop_rv The loop above which the new loop is created
|
|
* \return The new loop created
|
|
*/
|
|
virtual LoopRV AddUnitLoop(const LoopRV& loop_rv) = 0;
|
|
/******** Schedule: Manipulate ForKind ********/
|
|
/*!
|
|
* \brief Parallelize the input loop. It requires:
|
|
* 1) The scope block that the loop is in should have stage-pipeline property
|
|
* 2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
|
|
* bindings
|
|
* 3) For each block under the loop, the loop can only be contained in data-parallel block iters'
|
|
* bindings
|
|
* \param loop_rv The loop to be parallelized
|
|
*/
|
|
virtual void Parallel(const LoopRV& loop_rv) = 0;
|
|
/*!
|
|
* \brief Vectorize the input loop. It requires:
|
|
* 1) The scope block that the loop is in should have stage-pipeline property
|
|
* 2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
|
|
* bindings
|
|
* 3) For each block under the loop, the loop can only be contained in data-parallel block iters'
|
|
* bindings
|
|
* \param loop_rv The loop to be vectorized
|
|
*/
|
|
virtual void Vectorize(const LoopRV& loop_rv) = 0;
|
|
/*!
|
|
* \brief Bind the input loop to the given thread axis. It requires:
|
|
* 1) The scope block that the loop is in should have stage-pipeline property
|
|
* 2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
|
|
* bindings
|
|
* 3) For each block under the loop, if the thread axis starts with "threadIdx`, the loop can only
|
|
* be contained in data-parallel block iter and reduction block iters' bindings. Otherwise the
|
|
* loop can only be contained in data-parallel block iters' bindings
|
|
* \param loop_rv The loop to be bound to the thread axis
|
|
* \param thread_axis The thread axis to be bound to the loop
|
|
*/
|
|
virtual void Bind(const LoopRV& loop_rv, const ffi::String& thread_axis) = 0;
|
|
/*!
|
|
* \brief Unroll the input loop. It requires nothing
|
|
* \param loop_rv The loop to be unrolled
|
|
*/
|
|
virtual void Unroll(const LoopRV& loop_rv) = 0;
|
|
/******** Schedule: Insert cache stages ********/
|
|
/*!
|
|
* \brief Create a block that reads a buffer region into a read cache. It requires:
|
|
* 1) There is at most one block who writes the buffer in the scope.
|
|
* 2) The scope block have stage-pipeline property.
|
|
* \param block_rv The consumer block of the target buffer.
|
|
* \param read_buffer_index The index of the buffer in block's read region.
|
|
* \param storage_scope The target storage scope.
|
|
* \param consumer_blocks An optional list of consumers of the cache to rewrite.
|
|
* \return The cache stage block.
|
|
*/
|
|
virtual SBlockRV CacheRead(const SBlockRV& block_rv, int read_buffer_index,
|
|
const ffi::String& storage_scope,
|
|
const ffi::Array<SBlockRV> consumer_blocks = {}) = 0;
|
|
/*!
|
|
* \brief Create a block that writes a buffer region into a write cache. It requires:
|
|
* 1) There is only one block who writes the target buffer.
|
|
* 2) The scope block have stage-pipeline property.
|
|
* \param block_rv The producer of the buffer
|
|
* \param write_buffer_index The index of the buffer in block's write region
|
|
* \param storage_scope The target storage scope
|
|
* \param consumer_blocks An optional list of consumers to read from cache directly.
|
|
* \return The cache stage block.
|
|
*/
|
|
virtual SBlockRV CacheWrite(const SBlockRV& block_rv, int write_buffer_index,
|
|
const ffi::String& storage_scope,
|
|
const ffi::Array<SBlockRV> consumer_blocks = {}) = 0;
|
|
/*!
|
|
* \brief Create a block that reads a buffer region into a read cache. It requires:
|
|
* 1) There is at most one block who writes the buffer in the scope.
|
|
* 2) The scope block have stage-pipeline property.
|
|
* Compared to cache read, the indices to access allocated cache buffer is customized by user.
|
|
* \param block_rv The consumer block of the target buffer.
|
|
* \param read_buffer_index The index of the buffer in block's read region.
|
|
* \param storage_scope The target storage scope.
|
|
* \param index_map User defined indices to access allocated cache buffer, maps from block iter
|
|
* vars.
|
|
* \return The cache stage block.
|
|
*/
|
|
virtual SBlockRV ReindexCacheRead(const SBlockRV& block_rv, int read_buffer_index,
|
|
const ffi::String& storage_scope,
|
|
const IndexMap& index_map) = 0;
|
|
/*!
|
|
* \brief Create a block that writes a buffer region into a write cache. It requires:
|
|
* 1) There is only one block who writes the target buffer.
|
|
* 2) The scope block have stage-pipeline property.
|
|
* Compared to cache write, the indices to access allocated cache buffer is customized by user.
|
|
* \param block_rv The producer of the buffer
|
|
* \param write_buffer_index The index of the buffer in block's write region
|
|
* \param storage_scope The target storage scope
|
|
* \param index_map User defined indices to access allocated cache buffer, maps from block iter
|
|
* vars.
|
|
* \return The cache stage block.
|
|
*/
|
|
virtual SBlockRV ReindexCacheWrite(const SBlockRV& block_rv, int write_buffer_index,
|
|
const ffi::String& storage_scope,
|
|
const IndexMap& index_map) = 0;
|
|
/*!
|
|
* \brief Create 2 blocks that read&write a buffer region into a read/write cache.
|
|
* It requires the target block both read & write the target buffer.
|
|
* \param block_rv The target block operates on the target buffer.
|
|
* \param read_buffer_index The index of the buffer in block's read region.
|
|
* \param storage_scope The target storage scope
|
|
* \return The cache stage blocks, cache read block together with cache write block.
|
|
*/
|
|
virtual ffi::Array<SBlockRV> CacheInplace(const SBlockRV& block_rv, int read_buffer_index,
|
|
const ffi::String& storage_scope) = 0;
|
|
/*!
|
|
* \brief Create a block to cache precomputed index for later use.
|
|
* if there is no index computation, keep unchanged.
|
|
* \param block_rv The target block
|
|
* \param storage_scope The storage scope of cached block
|
|
* \param cse_thresh The repeat threshold that determines a common sub expr
|
|
* \return The cache stage blocks.
|
|
*/
|
|
virtual ffi::Array<SBlockRV> CacheIndex(const SBlockRV& block_rv,
|
|
const ffi::String& storage_scope, int cse_thresh) = 0;
|
|
/*!
|
|
* \brief Create a block that read/write a buffer region into a read/write cache with reindexing.
|
|
* The layout of the cache will be the same as by the iterators of the block that reads/writes the
|
|
* buffer. It requires:
|
|
* 1) There is only one block who reads/writes the target buffer
|
|
* 2) There is only one buffer load/store of this buffer in the block
|
|
* \param block_rv The block operates on the target buffer.
|
|
* \param buffer_index The index of the buffer in block's read or write region.
|
|
* \param buffer_index_type The type of the buffer index, kRead or kWrite.
|
|
* \return The reindex stage block.
|
|
*/
|
|
virtual SBlockRV ReIndex(const SBlockRV& block_rv, int buffer_index,
|
|
BufferIndexType buffer_index_type) = 0;
|
|
/******** Schedule: Data movement ********/
|
|
virtual SBlockRV ReadAt(const LoopRV& loop_rv, const SBlockRV& block_rv, int read_buffer_index,
|
|
const ffi::String& storage_scope) = 0;
|
|
virtual SBlockRV WriteAt(const LoopRV& loop_rv, const SBlockRV& block_rv, int write_buffer_index,
|
|
const ffi::String& storage_scope) = 0;
|
|
/******** Schedule: Compute location ********/
|
|
/*!
|
|
* \brief Move a producer block under the specific loop, and regenerate the
|
|
* loops induced by the block so that the buffer region produced by the producer block could
|
|
* cover those regions consumed by its consumer blocks under the given loop. It requires:
|
|
* 1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
|
|
* 2) The scope block has stage-pipeline property
|
|
* 3) The subtree of the scope block, where the given block is in, satisfies the compact dataflow
|
|
* condition. i.e. all the blocks in the scope block's subtree must be either complete block or
|
|
* reduction block
|
|
* 4) The block is not an output block with regard to the scope block, i.e. the buffers written by
|
|
* the block are allocated under the scope block
|
|
* 5) All the consumers of the block are under the given loop
|
|
* \param block_rv The block to be moved
|
|
* \param loop_rv The loop where the block to be moved under
|
|
* \param preserve_unit_loops Whether to keep the trivial loops whose extents are 1
|
|
* \param index The block index of the loop body subtree blocks:
|
|
* - `index = -1` means inserted into the last possible insertion point;
|
|
* - `index = -2` means inserted into the first possible insertion point;
|
|
* - Otherwise, `index` is a nonnegative number that indicates the insertion point
|
|
*/
|
|
virtual void ComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv, bool preserve_unit_loops,
|
|
int index = -1) = 0;
|
|
/*!
|
|
* \brief Move a consumer block under the specific loop, and regenerate the
|
|
* loops induced by the block so that the buffer region consumed by the consumer block could
|
|
* cover those regions produced by its producer blocks under the given loop. It requires:
|
|
* 1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
|
|
* 2) The scope block has stage-pipeline property
|
|
* 3) The subtree of the scope block, where the given block is in, satisfies the compact dataflow
|
|
* condition. i.e. all the blocks in the scope block's subtree must be either complete block or
|
|
* reduction block
|
|
* 4) All the producers of the block are under the given loop
|
|
*
|
|
* \param block_rv The block to be moved
|
|
* \param loop_rv The loop where the block to be moved under
|
|
* \param preserve_unit_loops Whether to keep the trivial loops whose extents are 1
|
|
* \param index The block index of the loop body subtree blocks:
|
|
* - `index = -1` means inserted into the last possible insertion point;
|
|
* - `index = -2` means inserted into the first possible insertion point;
|
|
* - Otherwise, `index` is a nonnegative number that indicates the insertion point
|
|
*/
|
|
virtual void ReverseComputeAt(const SBlockRV& block_rv, const LoopRV& loop_rv,
|
|
bool preserve_unit_loops, int index = -1) = 0;
|
|
/*!
|
|
* \brief Inline a block into its consumer(s). It requires:
|
|
* 1) The block is a complete non-root block, which only produces one buffer
|
|
* 2) The block must not be the only leaf in the scope.
|
|
* 3) The body of the block must be a BufferStore statement in the form of,
|
|
* A[i, j, k, ...] = ...
|
|
* where the indices of the LHS are all distinct atomic variables,
|
|
* and no variables other than those indexing variables are allowed in the statement.
|
|
* \param block The block to be inlined to its consumer(s)
|
|
*/
|
|
virtual void ComputeInline(const SBlockRV& block) = 0;
|
|
/*!
|
|
* \brief Inline a block into its only producer. It requires:
|
|
* 1) The block is a complete non-root block, which only produces and consumers one buffer
|
|
* 2) The block must not be the only leaf in the scope.
|
|
* 3) The only producer of the block is a read-after-write producer and a complete non-root block
|
|
* 4) The body of the block must be a BufferStore statement in the form of,
|
|
* B[f(i, j, k, ...)] = g(i, j, k, A[i, j, k, ...] ...)
|
|
* where the indices of each `BufferLoad` on the RHS are all distinct atomic variables,
|
|
* and no variables other than those indexing variables are allowed in the statement.
|
|
* \param block The block to be inlined to its producer
|
|
*/
|
|
virtual void ReverseComputeInline(const SBlockRV& block) = 0;
|
|
/*!
|
|
* \brief Fuse an epilogue block into a reduction block
|
|
* \param reduction_block The reduction block (e.g., matmul)
|
|
* \param epilogue_block The epilogue block to be fused (e.g., bias add)
|
|
*/
|
|
virtual void FuseReductionEpilogue(const SBlockRV& reduction_block,
|
|
const SBlockRV& epilogue_block) = 0;
|
|
/******** Schedule: Reduction ********/
|
|
/*!
|
|
* \brief Decompose a reduction block into two separate blocks.
|
|
* a) The init block, which is translated from the init statement of the reduction block;
|
|
* b) The update block, which is the original block without init statement.
|
|
*
|
|
* The init block is inserted right before the given loop.
|
|
*
|
|
* The schedule primitive requires:
|
|
* 1) The input block is a reduction block.
|
|
* 2) The input loop is the ancestor of the block.
|
|
* 3) The input loop is not lower than all the loops related to reduce block var.
|
|
* \param block_rv The reduction block to be decomposed
|
|
* \param loop_rv The loop above which the init block is inserted before.
|
|
* \return The init block
|
|
*/
|
|
virtual SBlockRV DecomposeReduction(const SBlockRV& block_rv, const LoopRV& loop_rv) = 0;
|
|
/*!
|
|
* \brief Factorize an associative reduction block by the specified loop.
|
|
* \details An associative reduction cannot be parallelized directly,
|
|
* because it leads to potential race condition during accumulation.
|
|
* Alternatively, the reduction could be factorized on a loop with the following steps:
|
|
* - Step 1: evenly slice the reduction into `n` separate chunks, where `n` is the loop extent
|
|
* - Step 2: compute the chunks separately and write the result into `n` intermediate buffers;
|
|
* - Step 3: accumulate the `n` separate buffer into the result buffer.
|
|
* Note that the Step 2 above introduces opportunities for parallelization.
|
|
* RFactor is a schedule primitive that implements the transformation described above.
|
|
* \param loop_rv The loop outside block we want to do rfactor
|
|
* \param factor_axis The position where the new dimension is placed in the new introduced rfactor
|
|
* buffer. Suppose the original reduction block writes to buffer `B` with
|
|
* ndim(B) dimensions, then `factor_axis` should be in range `[-ndim(B) - 1,
|
|
* ndim(B)]`, and the negative index will be normalized to a non-negative one
|
|
* \return The rfactor block
|
|
*/
|
|
virtual SBlockRV RFactor(const LoopRV& loop_rv, int factor_axis) = 0;
|
|
/******** Schedule: SBlock annotation ********/
|
|
/*!
|
|
* \brief Set alignment requirement for specific dimension such that
|
|
* stride[axis] == k * factor + offset for some k. This is useful to set memory layout for
|
|
* more friendly memory access pattern. For example, we can set alignment to be factor=2,
|
|
* offset=1 to avoid bank conflict for thread access on higher dimension in GPU shared
|
|
* memory.
|
|
* \param block_rv The producer block of the buffer
|
|
* \param buffer_index The index of the buffer in block's write region
|
|
* \param axis The dimension to be specified for alignment
|
|
* \param factor The factor multiple of alignment
|
|
* \param offset The required offset factor
|
|
*/
|
|
virtual void StorageAlign(const SBlockRV& block_rv, int buffer_index, int axis, int factor,
|
|
int offset) = 0;
|
|
/*!
|
|
* \brief Set the storage scope of a buffer, where the buffer is specified by a block and a
|
|
* write-index
|
|
* \param block_rv The producer block of the buffer
|
|
* \param buffer_index The index of the buffer in block's write region
|
|
* \param storage_scope The storage scope to be set
|
|
*/
|
|
virtual void SetScope(const SBlockRV& block_rv, int buffer_index,
|
|
const ffi::String& storage_scope) = 0;
|
|
/*!
|
|
* \brief Set the data type of a buffer, where the buffer is specified by a block and a
|
|
* write-index
|
|
* \note This schedule primitive is unsafe and may change correctness of program because of
|
|
* type conversion, please use with caution.
|
|
* \param block_rv The producer block of the buffer
|
|
* \param buffer_index the index of the buffer in block's write region
|
|
* \param dtype The data type to be set
|
|
*/
|
|
virtual void UnsafeSetDType(const SBlockRV& block_rv, int buffer_index,
|
|
const ffi::String& dtype) = 0;
|
|
/******** Schedule: Blockize & Tensorize ********/
|
|
/*!
|
|
* \brief Convert the subtree rooted at a specific loop into a block.
|
|
* \param loop_rv the root of the subtree
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
* \return the new block
|
|
*/
|
|
virtual SBlockRV Blockize(const LoopRV& loop_rv, bool preserve_unit_iters = true) = 0;
|
|
/*!
|
|
* \brief Convert specified blocks into a nested block.
|
|
* \param blocks the specified block to construct the new block
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
* \return the new block
|
|
*/
|
|
virtual SBlockRV Blockize(const ffi::Array<SBlockRV>& blocks,
|
|
bool preserve_unit_iters = true) = 0;
|
|
/*!
|
|
* \brief Tensorize the computation enclosed by loop with the tensor intrin.
|
|
* \param loop_rv The loop to be tensorized
|
|
* \param intrin Name of the tensor intrinsic
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
*/
|
|
virtual void Tensorize(const LoopRV& loop_rv, const ffi::String& intrin,
|
|
bool preserve_unit_iters = true) = 0;
|
|
/*!
|
|
* \brief Tensorize the computation enclosed by loop with the tensor intrin.
|
|
* \param block_rv The block to be tensorized
|
|
* \param intrin Name of the tensor intrinsic
|
|
* \param preserve_unit_iters Whether or not to preserve unit iterators in block bindings
|
|
*/
|
|
virtual void Tensorize(const SBlockRV& block_rv, const ffi::String& intrin,
|
|
bool preserve_unit_iters = true) = 0;
|
|
|
|
/******** Schedule: Annotation ********/
|
|
/*!
|
|
* \brief Annotate a loop with a key value pair
|
|
* \param loop_rv The loop to be annotated
|
|
* \param ann_key The annotation key
|
|
* \param ann_val The annotation value, a string or a ExprRV
|
|
*/
|
|
virtual void Annotate(const LoopRV& loop_rv, const ffi::String& ann_key, const Any& ann_val) = 0;
|
|
/*!
|
|
* \brief Annotate a block with a key value pair
|
|
* \param block_rv The block to be annotated
|
|
* \param ann_key The annotation key
|
|
* \param ann_val The annotation value, a string or a ExprRV
|
|
*/
|
|
virtual void Annotate(const SBlockRV& block_rv, const ffi::String& ann_key,
|
|
const Any& ann_val) = 0;
|
|
/*!
|
|
* \brief Unannotate a loop's annotation with key ann_key
|
|
* \param loop_rv The loop to be unannotated
|
|
* \param ann_key The annotation key
|
|
*/
|
|
virtual void Unannotate(const LoopRV& loop_rv, const ffi::String& ann_key) = 0;
|
|
/*!
|
|
* \brief Unannotate a block's annotation with key ann_key
|
|
* \param block_rv The block to be unannotated
|
|
* \param ann_key The annotation key
|
|
*/
|
|
virtual void Unannotate(const SBlockRV& block_rv, const ffi::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 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. If the schedule contains a producer block
|
|
* for the specified buffer, the pad value will be written as
|
|
* part of the producer block if possible, or after the producer
|
|
* block otherwise. Otherwise, if the buffer is an input, will
|
|
* insert an annotation block to state that the padding contains
|
|
* the known value.
|
|
*
|
|
* Note: If applied to an input buffer, the calling scope is
|
|
* responsible for ensuring that the pad_value is present.
|
|
* 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 SBlockRV& block_rv, int buffer_index,
|
|
BufferIndexType buffer_index_type, const IndexMap& index_map,
|
|
const ffi::Optional<IndexMap>& pad_value = std::nullopt,
|
|
bool assume_injective_transform = false) = 0;
|
|
|
|
/*!
|
|
* \brief Apply a transformation represented by IndexMap to block
|
|
* \details The block iters and the block body are transformed by the given index_map.
|
|
* Outer loops corresponding to each new block iter are regenerated.
|
|
* The index_map is required to be bijective affine since we need its inverse mapping.
|
|
* \param block_rv The block to be transformed
|
|
* \param index_map The transformation to apply.
|
|
*/
|
|
virtual void TransformBlockLayout(const SBlockRV& block_rv, const IndexMap& index_map) = 0;
|
|
|
|
/*!
|
|
* \brief Set the axis separator of a buffer, where the buffer is specified by a block and a read
|
|
* or write index
|
|
* \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 buffer_index_type The type of the buffer index, kRead or kWrite.
|
|
* \param axis_separators The axis separator of the buffer
|
|
*/
|
|
virtual void SetAxisSeparator(const SBlockRV& block_rv, int buffer_index,
|
|
BufferIndexType buffer_index_type,
|
|
const ffi::Array<IntImm>& axis_separators) = 0;
|
|
|
|
/******** Schedule: Padding ********/
|
|
/*!
|
|
* \brief Decompose a padding block into a block filling const pad values and a block
|
|
* writing in-bound values.
|
|
* \param block_rv The block that match the padding pattern.
|
|
* \param loop_rv The loop above which the const filling block is inserted before.
|
|
* \return The const pad value filling block.
|
|
*/
|
|
virtual SBlockRV DecomposePadding(const SBlockRV& block_rv, const LoopRV& loop_rv) = 0;
|
|
|
|
/*!
|
|
* \brief Pad the computation of Einsum.
|
|
* \param block_rv The block that matches the Einsum pattern.
|
|
* \param padding The padding for each block iter.
|
|
* \details This schedule primitives identifies the Einsum pattern in the block body, and find its
|
|
* producer blocks. It then pads the computation of the Einsum pattern and its producer blocks.
|
|
* The output buffer and the producer buffer is resized according to the padding size. It requires
|
|
* the output buffer and the producer buffer to be allocated inside the PrimFunc.
|
|
*
|
|
* The padding is a list of non-negative integers, each element corresponds to the padding for
|
|
* each block iter in the order of block iters. The block and its producer blocks should have
|
|
* trivial bindings, i.e. each block iter is bound to a single loop variable. After padding, the
|
|
* block iter extent and the corresponding outer loop is extended by the padding size.
|
|
*
|
|
* The size of the producer buffers are infered from the padding size of the Einsum computation.
|
|
* The producer buffers are padded by the initial value of the corresponding reduction.
|
|
*/
|
|
virtual void PadEinsum(const SBlockRV& block_rv, const ffi::Array<int64_t>& padding) = 0;
|
|
|
|
/******** Schedule: Buffer transformation ********/
|
|
/*!
|
|
* \brief Compute the target buffer via rolling buffering.
|
|
* \details This primitive selects the outermost rollable axis with a positive bound overlap that
|
|
* appears in the block's ancestor loops as `rolling axis`, fold and circularize the buffer along
|
|
* the rolling dimension, append block predicate to avoid recomputing overlapping elements.
|
|
* It requires:
|
|
* 1) The buffer to be an intermediate buffer defined via `alloc_buffer`.
|
|
* 2) The LCA of the producer and consumer of the buffer is a for loop, typically,
|
|
* the producer and consumer of the buffer are cascaded through compute_at.
|
|
* 3) The access region of the buffer has at least one dimension that contains
|
|
* a positive bound overlap.
|
|
* \param block_rv The producer block of the buffer.
|
|
* \param write_buffer_index The index of the buffer in block's write region.
|
|
*/
|
|
virtual void RollingBuffer(const SBlockRV& block_rv, int write_buffer_index) = 0;
|
|
|
|
/*!
|
|
* \brief Annotate the buffer access of a block
|
|
* \param block_rv The block to be annotated
|
|
* \param buffer_index The index of the buffer in block's read or write region
|
|
* \param buffer_index_type The type of the buffer index, kRead or kWrite.
|
|
* \param index_map The index map that defines the new read or write region
|
|
*/
|
|
virtual void AnnotateBufferAccess(const SBlockRV& 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 */
|
|
virtual void EnterPostproc() = 0;
|
|
|
|
/*!
|
|
* \brief Hide some buffer access in the given block.
|
|
* \param block_rv The block where we hide buffer access.
|
|
* \param buf_type The buffer type: read/write
|
|
* \param buf_index_array The array of buffer indices we hide access.
|
|
*/
|
|
virtual void UnsafeHideBufferAccess(const SBlockRV& block_rv, const ffi::String& buf_type,
|
|
const ffi::Array<IntImm>& buf_index_array) = 0;
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to ScheduleNode
|
|
*
|
|
* A schedule is a set of transformations that change the order of computation but
|
|
* preserve the semantics of computation. Some example of schedules:
|
|
* 1) Split a loop into two;
|
|
* 2) Reorder two loops;
|
|
* 3) Inline the computation of a specific buffer into its consumer
|
|
*
|
|
* The schedule class stores auxiliary information to schedule correctly and efficiently.
|
|
*
|
|
* Link to tutorial: https://tvm.apache.org/docs/tutorials/language/schedule_primitives.html
|
|
*
|
|
* \sa ScheduleNode
|
|
*/
|
|
class Schedule : public ffi::ObjectRef {
|
|
public:
|
|
/*!
|
|
* \brief Construct a concrete TensorIR schedule from an IRModule
|
|
* \param mod The IRModule to be scheduled
|
|
* \param seed The seed value for schedule's random state
|
|
* \param debug_mask Do extra correctness checking after the class creation
|
|
* and each time after calling the Replace method.
|
|
* \param error_render_level The level of error rendering
|
|
* \param enable_check Whether to enable some prequisite checks for schedule primitives, it's
|
|
* user's duty to guarantee the schedule correctness if we disable the checks.
|
|
* \return The concrete schedule created
|
|
* \sa ScheduleDebugMask
|
|
* \note The checks performed includes: 1) VerifySRefTree 2) VerifyCachedFlags
|
|
*/
|
|
TVM_DLL static Schedule Concrete(IRModule mod, LinearCongruentialEngine::TRandState seed,
|
|
int debug_mask, ScheduleErrorRenderLevel error_render_level,
|
|
bool enable_check = true);
|
|
/*!
|
|
* \brief Construct a traced concrete TensorIR schedule from an IRModule
|
|
* \param mod The IRModule to be scheduled
|
|
* \param seed The seed value for schedule's random state
|
|
* \param debug_mask Do extra correctness checking after the class creation
|
|
* and each time after calling the Replace method.
|
|
* \param error_render_level The level of error rendering
|
|
* \param enable_check Whether to enable prequisite checks for schedule primitives.
|
|
* \return The concrete schedule created
|
|
* \sa ScheduleDebugMask
|
|
* \note The checks performed include:
|
|
* 1) VerifySRefTree
|
|
* 2) VerifyCachedFlags
|
|
*/
|
|
TVM_DLL static Schedule Traced(IRModule mod, LinearCongruentialEngine::TRandState seed,
|
|
int debug_mask, ScheduleErrorRenderLevel error_render_level,
|
|
bool enable_check = true);
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Schedule, ffi::ObjectRef, ScheduleNode);
|
|
};
|
|
|
|
} // namespace s_tir
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_S_TIR_SCHEDULE_SCHEDULE_H_
|