/* * 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_