wehub-resource-sync
375 lines
11 KiB
C++
375 lines
11 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/transform.h
|
|
* \brief S-TIR specific transformation passes.
|
|
*/
|
|
#ifndef TVM_S_TIR_TRANSFORM_H_
|
|
#define TVM_S_TIR_TRANSFORM_H_
|
|
|
|
#include <tvm/ir/transform.h>
|
|
#include <tvm/target/target.h>
|
|
#include <tvm/tirx/transform.h>
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace tvm {
|
|
namespace s_tir {
|
|
|
|
/*!
|
|
* \brief Renew the definition nodes for a TIR, including Var, Buffer and IterVar.
|
|
* This pass works as a simple DeepCopy to duplicate a function with different Vars and
|
|
* Buffers but the same behavior
|
|
* \param func The input PrimFunc.
|
|
* \return The renewed func.
|
|
*/
|
|
TVM_DLL tirx::PrimFunc RenewDefs(const tirx::PrimFunc& func);
|
|
|
|
namespace transform {
|
|
|
|
using tirx::transform::CreatePrimFuncPass;
|
|
using tvm::transform::Pass;
|
|
using tvm::transform::PassContext;
|
|
|
|
/*!
|
|
* \brief Canonicalize loop to start from zero .
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass CanonicalizeLoop();
|
|
|
|
/*!
|
|
* \brief Lower cross-thread reduction from thread
|
|
* bindings to intrinsic function calls.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerCrossThreadReduction();
|
|
|
|
/*!
|
|
* \brief Lower block init stmt into IfThenElse stmts
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerInitBlock();
|
|
|
|
/*!
|
|
* \brief Locate the buffer allocation to the exact position (usually is
|
|
* the lca of buffer access). This pass will inject opaque block
|
|
* with alloc_buffers at the allocation site.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass PlanAndUpdateBufferAllocationLocation();
|
|
|
|
/*!
|
|
* \brief Substitute all the block vars with the PrimExprs they are bound to, indicated by the
|
|
* corresponding iter_values in BlockRealize, for opaque blocks by removing all
|
|
*. the iter_values in BlockRealize and iter_vars in Block.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass ConvertBlocksToOpaque();
|
|
|
|
/*!
|
|
* \brief Lift the same thread bindings to their LCA loops
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LiftThreadBinding();
|
|
|
|
/*!
|
|
* \brief Compact the buffer access region by removing the buffer regions that are not accessed,
|
|
* i.e. narrowing the buffer shape and adjust the access region if necessary.
|
|
*
|
|
* Before narrowing, `B` is a `[16, 16]` buffer, but only a skinny vector `B[i, 0:16]` is accessed.
|
|
*
|
|
* \code
|
|
*
|
|
* for i in range(0, 16):
|
|
* with T.sblock():
|
|
* B = T.alloc_buffer(16, 16)
|
|
* for j in range(0, 16):
|
|
* B[i, j] = A[i, j] + 1
|
|
* for j in range(0, 16):
|
|
* C[i, j] = B[i, j] + 1
|
|
*
|
|
* \endcode
|
|
*
|
|
* This pass narrows the buffer shape and adjust its accessed region accordingly.
|
|
* In this particular case, because only a `1 * 16` vector of `B` is accessed,
|
|
* the pass narrows `B` to shape `[1, 16]`, and changes the access to `B[i, j]` to `B[0, j]`.
|
|
*
|
|
* \code
|
|
*
|
|
* for i in range(0, 16):
|
|
* with T.sblock():
|
|
* B = T.alloc_buffer(1, 16)
|
|
* for j in range(0, 16):
|
|
* B[0, j] = A[i, j] + 1
|
|
* for j in range(0, 16):
|
|
* C[i, j] = B[0, j] + 1
|
|
*
|
|
* \endcode
|
|
*
|
|
* \param is_strict ensure the compacted shape always smaller than the original shape.
|
|
* otherwise it allows to grow the shape to match actual accessed buffer regions.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass CompactBufferAllocation(bool is_strict = true);
|
|
|
|
/*!
|
|
* \brief Remove match buffers inside the block. Also, it will validate the binding.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerMatchBuffer();
|
|
|
|
/*!
|
|
* \brief Inject permuted layout for shared memory.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InjectPermutedLayout();
|
|
|
|
/*!
|
|
* \brief Transform Mma scope (m16n8k8.matrixA/B/C) to local scope with layout transformation.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass TransformMmaBufferLayout();
|
|
|
|
/*!
|
|
* \brief Remove the block to ensure that the TIR can not be scheduled again.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerOpaqueBlock();
|
|
|
|
/*!
|
|
* \brief Unify all the thread bindings for "blockIdx.x/y/z", "threadIdx.x/y/z", and
|
|
* "vthread.x/y/z". Before the unification, two vars that are bound to a thread axis (e.g.,
|
|
* "threadIdx.x") use different IterVars and variables in their AttrStmts. After the
|
|
* unification, we use a consolidated IterVar and a variable for them.
|
|
* \return The pass.
|
|
* \note `vthread` is a legacy behavior that will be deprecated, though thread bindings of `vthread`
|
|
* are still also unified in this pass. Please use `vthread.x`, `vthread.y` and `vthread.z`
|
|
* instead.
|
|
*/
|
|
TVM_DLL Pass UnifyThreadBinding();
|
|
|
|
/*!
|
|
* \brief This pass transforms annotated loops into pipelined ones where producers and consumers
|
|
* are overlapped with the information provided in loop annotations, which enables optimization
|
|
* techniques like prefetching and pipeline parallelism.
|
|
*
|
|
* The pipeline scope consists of the direct children of the annotated loop (ignoring SBlockRealize,
|
|
* SBlock, SeqStmt), and the number of children is denoted by `n` in the documentation.
|
|
*
|
|
* The following annotations are used to guide the loop transformation:
|
|
*
|
|
* 1) Loop annotation `software_pipeline_stage` defines the pipeline stage.
|
|
* An array of `n` integers, and each element should be in range [0, max_stage],
|
|
* where max_stage is the maximum (inclusive) stage.
|
|
* 2) Loop annotation `software_pipeline_order` defines the pipeline order.
|
|
* An array of `n` integers, a permutation of [0, 1, ..., num_components - 1];
|
|
* 3) SBlock annotation `double_buffer_scope` controls certain buffer sizes to allow decoupling of
|
|
* read/write dependency. It's an integer index of the write regions of the block.
|
|
*
|
|
* Every annotated loop is transformed into a loop with three blocks as its direct children:
|
|
*
|
|
* 1) Prologue block, where components whose stage is less than `max_stage` is executed;
|
|
*
|
|
* 2) Body block, where all the components are executed;
|
|
*
|
|
* 3) Epilogue block, where only components whose stage is greater than 0 will be executed.
|
|
* The execution order is controlled by the annotation `software_pipeline_order`,
|
|
* and thus could be different than the original order.
|
|
*
|
|
* Note: For nested software pipelines, the inner software pipeline will be generated first,
|
|
* which may affect the number of the direct children of the outer loop.
|
|
* In this case, the annotations for the outer software
|
|
* pipeline should include the result of the inner software pipeline,
|
|
* which is the three blocks as discussed above.
|
|
*
|
|
* \return The IR transform pass.
|
|
*/
|
|
TVM_DLL Pass InjectSoftwarePipeline();
|
|
|
|
/*!
|
|
* \brief Automatically do memory optimizations for auto copy blocks
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerAutoCopy();
|
|
|
|
/*!
|
|
* \brief Add the explicit local stage for the shared memory access on GPU.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass ManifestSharedMemoryLocalStage();
|
|
|
|
/*! \brief Annotate irregular loop mark. */
|
|
TVM_DLL Pass AnnotateIrregularLoop();
|
|
|
|
/*!
|
|
* \brief partition loops in the stmt.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LoopPartition();
|
|
|
|
/*!
|
|
* \brief Inject virtual thread loops.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InjectVirtualThread();
|
|
|
|
/*!
|
|
* \brief Inject double buffer statements.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InjectDoubleBuffer();
|
|
|
|
/*!
|
|
* \brief Hoist loop-invariant IfThenElse nodes to
|
|
* outside the eligible loops.
|
|
*
|
|
* \param variant The variant of the pass.
|
|
* variant can have any one of following values ["basic", ""(Default)].
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass HoistIfThenElse(tvm::ffi::String variant = "");
|
|
|
|
/*!
|
|
* \brief Hoist loop-invariant expressions to outside the eligible loops.
|
|
*
|
|
* Can hoist conditionals used in IfThenElse statements and
|
|
* expressions, bindings of variables in Let statements and
|
|
* expressions, or boolean expressions, configurable to enable/disable
|
|
* each hoistable type.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass HoistExpression();
|
|
|
|
/*!
|
|
* \brief Renormalize the split pattern from floordiv(floormod()) to floormod(floordiv()).
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RenormalizeSplitPattern();
|
|
|
|
/*!
|
|
* \brief Detect and rewrite unsafe select that contains memory access.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RewriteUnsafeSelect();
|
|
|
|
/*!
|
|
* \brief Instruments bound checkers.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InstrumentBoundCheckers();
|
|
|
|
/*!
|
|
* \brief Rewrite global to local memory copy on CUDA with ldg32 instruction.
|
|
* \param enable_inject Whether to enable injection.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InjectPTXLDG32(bool enable_inject = true);
|
|
|
|
/*!
|
|
* \brief Insert intrinsic calls to instrument function and loop level profiling.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InstrumentProfileIntrinsics();
|
|
|
|
/*!
|
|
* \brief Lower VTCM allocations.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerVtcmAlloc();
|
|
|
|
/*!
|
|
* \brief Insert sync between parallel read/write of shared buffers.
|
|
* \param storage_scope The storage scope considered.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass ThreadSync(tvm::ffi::String storage_scope);
|
|
|
|
/*!
|
|
* \brief Infer the TensorCore fragment information using tensor intrinsics.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InferFragment();
|
|
|
|
/*!
|
|
* \brief Lower cross thread allreduce.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerThreadAllreduce();
|
|
|
|
/*!
|
|
* \brief Lower Async TIR primitives to DMA copy and wait builtins.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerAsyncDMA();
|
|
|
|
/*!
|
|
* \brief Rewrite global to shared memory copy on CUDA with asynchronous copy.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InjectPTXAsyncCopy();
|
|
|
|
/*!
|
|
* \brief Merge multiple TIR-level shared memory allocations into one.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass MergeSharedMemoryAllocations();
|
|
|
|
/*!
|
|
* \brief Set default thread bindings for GPU PrimFuncs.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass DefaultGPUSchedule();
|
|
|
|
/*!
|
|
* \brief Remove weight layout rewrite block before benchmark.
|
|
* \param skip_tensor_rewrite Whether to skip tensor rewrite.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RemoveWeightLayoutRewriteBlock(bool skip_tensor_rewrite = false);
|
|
|
|
/*!
|
|
* \brief Remove stores of tirx::builtin::undef.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RemoveStoreUndef();
|
|
|
|
/*!
|
|
* \brief Decorate all the function's body as device function.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass DecorateDeviceScope();
|
|
|
|
/*!
|
|
* \brief Eliminate branches by leveraging buffer assumptions (T.assume).
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass UseAssumeToReduceBranches();
|
|
|
|
} // namespace transform
|
|
} // namespace s_tir
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_S_TIR_TRANSFORM_H_
|