244 lines
9.1 KiB
C++
244 lines
9.1 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.
|
|
*/
|
|
/*!
|
|
* \file tvm/s_tir/stmt.h
|
|
* \brief S-TIR (Schedulable TIR) statement attribute declarations.
|
|
*
|
|
* This file contains attribute keys that are specific to the schedulable TIR
|
|
* (S-TIR) layer, including meta_schedule annotations and schedule primitive /
|
|
* SBlock annotations.
|
|
*/
|
|
#ifndef TVM_S_TIR_STMT_H_
|
|
#define TVM_S_TIR_STMT_H_
|
|
|
|
namespace tvm {
|
|
namespace s_tir {
|
|
namespace attr {
|
|
|
|
/*!
|
|
* \brief Annotations for invoking and synchronizing asynchronous operations.
|
|
*/
|
|
constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
|
|
constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
|
|
constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
|
|
|
|
/*!
|
|
* \brief Mark that the attached statement runs asynchronously.
|
|
*/
|
|
constexpr const char* async_scope = "async_scope";
|
|
|
|
/*!
|
|
* \brief Marks production of double buffer data
|
|
*/
|
|
constexpr const char* double_buffer_scope = "double_buffer_scope";
|
|
|
|
/*!
|
|
* \brief Marks region used by double buffer write
|
|
*/
|
|
constexpr const char* double_buffer_write = "double_buffer_write";
|
|
|
|
/*!
|
|
* \brief Mark that the shape of TensorCore fragment
|
|
*/
|
|
constexpr const char* fragment_shape = "fragment_shape";
|
|
|
|
/*!
|
|
* \brief Mark that the layout of TensorCore fragment
|
|
*/
|
|
constexpr const char* fragment_layout = "fragment_layout";
|
|
|
|
/*!
|
|
* \brief Mark that the loop should be partitioned.
|
|
*/
|
|
constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
|
|
|
|
/*! \brief Mark of reduce scope */
|
|
constexpr const char* reduce_scope = "reduce_scope";
|
|
|
|
/*! \brief Mark launching of a virtual thread. */
|
|
constexpr const char* virtual_thread = "virtual_thread";
|
|
|
|
// -----------------------------------------------------------------------
|
|
// meta_schedule annotations
|
|
// -----------------------------------------------------------------------
|
|
|
|
/*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */
|
|
constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure";
|
|
|
|
/*!
|
|
* \brief Mark that the loop should be further skip and bound to environment threads to enable
|
|
* cooperative fetching.
|
|
*/
|
|
constexpr const char* meta_schedule_cooperative_fetch = "meta_schedule.cooperative_fetch";
|
|
|
|
/*! \brief The allowed range of thread extent in thread bindings */
|
|
constexpr const char* meta_schedule_thread_extent_low_inclusive =
|
|
"meta_schedule.thread_extent_low_inclusive";
|
|
|
|
/*! \brief The allowed range of thread extent in thread bindings */
|
|
constexpr const char* meta_schedule_thread_extent_high_inclusive =
|
|
"meta_schedule.thread_extent_high_inclusive";
|
|
|
|
/*! \brief Mark the block whose producer needs to be applied by rule Random-Compute-Location */
|
|
constexpr const char* meta_schedule_random_compute_producer =
|
|
"meta_schedule.random_compute_producer";
|
|
|
|
/*! \brief Mark auto-parallel setting on the block. */
|
|
constexpr const char* meta_schedule_parallel = "meta_schedule.parallel";
|
|
|
|
/*! \brief Mark auto-vectorize setting on the block. */
|
|
constexpr const char* meta_schedule_vectorize = "meta_schedule.vectorize";
|
|
|
|
/*! \brief Mark auto-unroll setting on the block. */
|
|
constexpr const char* meta_schedule_unroll_explicit = "meta_schedule.unroll_explicit";
|
|
|
|
/*! \brief Mark auto-unroll setting on the block. */
|
|
constexpr const char* meta_schedule_unroll_implicit = "meta_schedule.unroll_implicit";
|
|
|
|
/*! \brief Mark that a block should be further rewritten using tensorization. */
|
|
constexpr const char* meta_schedule_auto_tensorize = "meta_schedule.auto_tensorize";
|
|
|
|
/*! \brief Mark that a block is a preprocessor block for layout rewrite. */
|
|
constexpr const char* meta_schedule_layout_rewrite_preproc = "meta_schedule.layout_rewrite_preproc";
|
|
|
|
/*!
|
|
* \brief Mark that the init statement of a block should be further rewritten using tensorization.
|
|
*/
|
|
constexpr const char* meta_schedule_auto_tensorize_init = "meta_schedule.auto_tensorize_init";
|
|
|
|
/*! \brief Mark that tensor core is enabled in the PrimExpr */
|
|
constexpr const char* meta_schedule_tensor_core_enabled = "meta_schedule.tensor_core_enabled";
|
|
|
|
/*!
|
|
* \brief Mark a block as generated by cache_read or cache_write block.
|
|
* 0 means cache_read; 1 means cache_write.
|
|
* \sa meta_schedule_cache_type_read
|
|
* \sa meta_schedule_cache_type_write
|
|
*/
|
|
constexpr const char* meta_schedule_cache_type = "meta_schedule.cache_type";
|
|
|
|
/*! \sa meta_schedule_cache_type */
|
|
constexpr const int meta_schedule_cache_type_read = 0;
|
|
|
|
/*! \sa meta_schedule_cache_type */
|
|
constexpr const int meta_schedule_cache_type_write = 1;
|
|
|
|
/*! \brief Mark that a block is disallowed in auto inline. */
|
|
constexpr const char* meta_schedule_inline_rule = "meta_schedule.inline_rule";
|
|
|
|
// -----------------------------------------------------------------------
|
|
// Schedule primitive / SBlock annotations
|
|
// -----------------------------------------------------------------------
|
|
|
|
/*!
|
|
* \brief Mark whether the script-completer need to fill in missing access region
|
|
* during script parsing.
|
|
* \note The result should be a integer mask with range [0, 4).
|
|
* if (mask & 1) the read region should be detected,
|
|
* if (mask & 2) the write region should be detected.
|
|
*/
|
|
constexpr const char* script_parsing_detect_access = "tirx.script_parsing_detect_access";
|
|
|
|
/*!
|
|
* \brief Mark that the block need to add predicate for block var bounds during lowering
|
|
*/
|
|
constexpr const char* require_block_var_bound_predicate = "require_bound_predicate";
|
|
|
|
/*! \brief Mark the stage of a statement in the software pipeline */
|
|
constexpr const char* software_pipeline_stage = "software_pipeline_stage";
|
|
|
|
/*! \brief Mark the order of a statement in the software pipeline */
|
|
constexpr const char* software_pipeline_order = "software_pipeline_order";
|
|
|
|
/*! \brief List stages in the software pipeline that should run asynchronously
|
|
* \note All statements in the provided stages are assumed to have asynchronous
|
|
* semantics (e.g. CUDA async global to shared memory copy).
|
|
*/
|
|
constexpr const char* software_pipeline_async_stages = "software_pipeline_async_stages";
|
|
|
|
/*! \brief Mark the buffers which is const access and can be transformed layout. */
|
|
constexpr const char* layout_free_buffers = "layout_free_buffers";
|
|
|
|
/*! \brief Mark the local stage for the shared memory access should be added. */
|
|
constexpr const char* manifest_shared_memory_local_stage =
|
|
"tirx.manifest_shared_memory_local_stage";
|
|
|
|
/*!
|
|
* \brief Mark alignment of buffer dimension
|
|
* stmt.node is Tensor
|
|
* stmt.value is tvm_tuple(dim, align, offset)
|
|
* This gives hint to require stride of dim to be k * align + offset.
|
|
*/
|
|
constexpr const char* buffer_dim_align = "buffer_dim_align";
|
|
|
|
/*! \brief Mark that a block has an explicitly specified read region.
|
|
* This is used to override the default read region inference in TIR.
|
|
*/
|
|
constexpr const char* explicit_read_region = "explicit_read_region";
|
|
|
|
/*! \brief Mark that a block has an explicitly specified write region.
|
|
* This is used to override the default write region inference in TIR.
|
|
*/
|
|
constexpr const char* explicit_write_region = "explicit_write_region";
|
|
|
|
/*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */
|
|
constexpr const char* irregular_loop_mark = "irregular_loop_mark";
|
|
|
|
/*! \brief Mark auto copy for memhammer */
|
|
constexpr const char* auto_copy = "auto_copy";
|
|
|
|
/*! \brief Mark local stage constraint on data copy */
|
|
constexpr const char* local_stage = "local_stage";
|
|
|
|
/*! \brief Mark vectorization length constraint on block */
|
|
constexpr const char* vector_bytes = "vector_bytes";
|
|
|
|
/*!
|
|
* \brief Mark that a block is executed by a warp. This implies the extend of threadIdx.x is
|
|
* warp size.
|
|
*/
|
|
constexpr const char* warp_execution = "warp_execution";
|
|
|
|
/*!
|
|
* \brief Marks the layout transforms to be used for a tensor.
|
|
*
|
|
* Only applies to a DataProducer, as it should be made part of the
|
|
* PrimFunc attributes for TIR.
|
|
*/
|
|
constexpr const char* layout_transforms = "layout_transforms";
|
|
|
|
/*!
|
|
* \brief Marks the physical axis separators
|
|
*
|
|
* Only applies to a DataProducer, as it should be made part of the
|
|
* Buffer definition in a PrimFunc. See `BufferNode::axis_separators`
|
|
* for more details.
|
|
*/
|
|
constexpr const char* axis_separators = "axis_separators";
|
|
|
|
/*!
|
|
* \brief Mark that the kernel is hand threaded and doesn't need syncs inserted
|
|
*/
|
|
constexpr const char* hand_threaded = "hand_threaded";
|
|
|
|
} // namespace attr
|
|
} // namespace s_tir
|
|
} // namespace tvm
|
|
#endif // TVM_S_TIR_STMT_H_
|