wehub-resource-sync
346 lines
9.4 KiB
C++
346 lines
9.4 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/tirx/transform.h
|
|
* \brief TIR specific transformation passes.
|
|
*/
|
|
#ifndef TVM_TIR_TRANSFORM_H_
|
|
#define TVM_TIR_TRANSFORM_H_
|
|
|
|
#include <tvm/ir/transform.h>
|
|
#include <tvm/target/target.h>
|
|
#include <tvm/tirx/expr.h>
|
|
#include <tvm/tirx/function.h>
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace tvm {
|
|
namespace tirx {
|
|
namespace transform {
|
|
|
|
using tvm::transform::CreateModulePass;
|
|
using tvm::transform::Pass;
|
|
using tvm::transform::PassContext;
|
|
using tvm::transform::PassContextNode;
|
|
using tvm::transform::PassInfo;
|
|
using tvm::transform::PassInfoNode;
|
|
using tvm::transform::PassNode;
|
|
using tvm::transform::Sequential;
|
|
|
|
/*
|
|
* \brief Create a function pass that optimizes PrimFuncs.
|
|
*
|
|
* \param pass_func The packed function that contains the optimization.
|
|
* \param opt_level The optimization level of the function pass.
|
|
* \param name The name of the function pass.
|
|
* \param required The list of the passes that the function pass is dependent on.
|
|
*
|
|
* \return The created function pass.
|
|
*/
|
|
TVM_DLL Pass CreatePrimFuncPass(std::function<PrimFunc(PrimFunc, IRModule, PassContext)> pass_func,
|
|
int opt_level, ffi::String name,
|
|
tvm::ffi::Array<ffi::String> required, bool traceable = false);
|
|
|
|
/*!
|
|
* \brief Lower vectorization loops.
|
|
*
|
|
* \param enable_vectorize Whether vectorization is enabled.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass VectorizeLoop(bool enable_vectorize = true);
|
|
|
|
/*!
|
|
* \brief Rewrite storage allocation pattern.
|
|
* Moves the allocation to outer most possible scope.
|
|
* Trying to share space between allocations to make
|
|
* a static allocation plan when possible.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass StorageRewrite();
|
|
|
|
/*!
|
|
* \brief unroll the constant loop marked by unroll.
|
|
* This pass also automatically attach pragma unroll tag to loops which meets the standard.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass UnrollLoop();
|
|
|
|
/*!
|
|
* \brief Remove No Op from the Stmt.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RemoveNoOp();
|
|
|
|
/*!
|
|
* \brief Run statement-level arithmetic simplifications on the TIR PrimFunc.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass StmtSimplify();
|
|
|
|
/*!
|
|
* \brief Convert an IRModule to be SSA form.
|
|
*
|
|
* This pass handles cases where the same tirx::Var appears in
|
|
* multiple functions within the same module. For example, after
|
|
* extracting a fragment from one function into another, where the
|
|
* same `tirx::Var` may be defined both as within the body of the
|
|
* original function, and as a parameter within the hoisted function.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass ConvertSSA();
|
|
|
|
/*!
|
|
* \brief Transform the high-level PrimFunc to a low-level version
|
|
* that can be used as an API function.
|
|
*
|
|
*
|
|
* The main task of this function is to create code to :
|
|
* - Map the values in the api_args to Var that is required by body.
|
|
* - Insert assertions to check type/value of the passed arguments.
|
|
*
|
|
* \note
|
|
* The function signature have two cases
|
|
*
|
|
* let num_packed_args = len(api_args);
|
|
*
|
|
* if num_packed_args is zero:
|
|
* f()
|
|
*
|
|
* if num_packed_args is not zero:
|
|
* f(void *, TVMFFIAny* packed_args, int num_packed_args,
|
|
* api_arg_k, api_arg_k+1, ... api_arg_n,
|
|
* TVMFFIAny* out_ret_val)
|
|
*
|
|
* where n == len(api_args), k == num_packed_args
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass MakePackedAPI();
|
|
|
|
/*!
|
|
* \brief Remap the thread axis
|
|
*
|
|
* This can be used to get equivalent program which uses
|
|
* threadIdx.y in place of threadIdx.x by passing
|
|
* {"threadIdx.x": thread_axis("threadIdx.y")}
|
|
*
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass RemapThreadAxis(ffi::Map<ffi::String, IterVar> axis_map);
|
|
|
|
/*!
|
|
* \brief Annotate, split, and lower host/device functions.
|
|
*
|
|
* This pass first annotates device regions within host functions,
|
|
* then splits them into host and device-side PrimFuncs, and finally
|
|
* lowers host-to-device calls into the device kernel launch ABI.
|
|
*
|
|
* The resulting host-side function will keep the same
|
|
* `tvm::attr::kTarget` attribute (e.g. `T.target("cuda",
|
|
* host=T.target("llvm"))`). This ensures that `MakePackedAPI` knows
|
|
* which device type should be used for the input buffers.
|
|
*
|
|
* The resulting device-side function will
|
|
* have the host stripped from its target attribute
|
|
* (e.g. `T.target("cuda")`).
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass SplitHostDevice();
|
|
|
|
/*!
|
|
* \brief skip assert stmt.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass SkipAssert();
|
|
|
|
/*!
|
|
* \brief This annotation is for nodes to be disabled for builtin lowering
|
|
*/
|
|
static constexpr const char* kDisableLowerTVMBuiltin = "disable_lower_builtin";
|
|
|
|
/*!
|
|
* \brief Lower builtin intrinsics.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerTVMBuiltin();
|
|
|
|
/*!
|
|
* \brief Lower the target specific function intrinsics in each of the function.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerIntrin();
|
|
|
|
/*!
|
|
* \brief Lower warp memory access to low-level device related function calls.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerWarpMemory();
|
|
|
|
/*!
|
|
* \brief Narrow down PrimExpr datatype in stmt to target_bits.
|
|
*
|
|
* \param target_bits The target bits
|
|
*
|
|
* \note Run this pass after storage flatten.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass NarrowDataType(int target_bits);
|
|
|
|
/*!
|
|
* \brief Force to narrow down indexing expressions and integer buffers to int32 dtype.
|
|
*
|
|
* \return The pass.
|
|
* \note This pass should not be used in default cases.
|
|
*/
|
|
TVM_DLL Pass ForceNarrowIndexToInt32();
|
|
|
|
/*!
|
|
* \brief Legalize bf16 compute Ops. Add a cast to fp32
|
|
* before Ops, then add a cast back to bf16.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass BF16ComputeLegalize();
|
|
|
|
/*!
|
|
* \brief Legalize fp8 compute Ops. Add a cast to fp16/fp32
|
|
* before Ops, then add a cast back to fp8.
|
|
* \param promote_dtype The data type used for type promotion, defaults to float16
|
|
* \note Must be run after BindTarget, as it relies on target attributes for PrimFuncs
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass FP8ComputeLegalize(ffi::String promote_dtype = "float16");
|
|
|
|
/*!
|
|
* \brief Legalize bf16 storage types to u16.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass BF16StorageLegalize();
|
|
|
|
/*!
|
|
* \brief Legalize fp8 storage types to u8.
|
|
* \note Must be run after BindTarget, as it relies on target attributes for PrimFuncs
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass FP8StorageLegalize();
|
|
|
|
/*!
|
|
* \brief Inline calls to private functions
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass InlinePrivateFunctions();
|
|
|
|
/*!
|
|
* \brief Rewrite the pointer content type of arguments,
|
|
* as well as Alloc internal to the function to use
|
|
* the most frequently accessed type for load/store
|
|
* to avoid pointer casting in backend when possible.
|
|
*
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass PointerValueTypeRewrite();
|
|
|
|
/*!
|
|
* \brief Flatten the multi-dimensional BufferLoad and BufferStore to single dimensional
|
|
* BufferLoad/BufferStore for the TIR not contains opaque block.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass FlattenBuffer();
|
|
|
|
/*!
|
|
* \brief Implements Common Subexpression Elimination (CSE) for TIR
|
|
* which introduces Bind statements for duplicated sub-expressions.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass CommonSubexprElim();
|
|
|
|
/*!
|
|
* \brief This is the unified static memory planner pass that will
|
|
* plan for memory intra- and inter- PrimFuncs together. The pass
|
|
* requires all the function to be PrimFuncs including the main.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass UnifiedStaticMemoryPlanner();
|
|
|
|
/*!
|
|
* \brief Annotate a PrimFunc with a given target.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass BindTarget(Target target);
|
|
|
|
/*!
|
|
* \brief Set a PrimFunc as the entry point if it is only function in IRModule.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass AnnotateEntryFunc();
|
|
|
|
/*!
|
|
* \brief Filter PrimFuncs with a given condition.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass Filter(ffi::TypedFunction<bool(PrimFunc)> fcond);
|
|
|
|
/*!
|
|
* \brief Lower TIRx op calls using registered op dispatchers for the given target.
|
|
*
|
|
* Also resolves ScopeIdDef declarations: gathers them at kernel scope, verifies
|
|
* consistency, extracts launch parameters, and emits Bind statements +
|
|
* thread_extent AttrStmts wrapping the dispatched body.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass TilePrimitiveDispatch();
|
|
|
|
/*!
|
|
* \brief Finalize TIRx lowering by applying layout rewriters and cleanup passes.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerTIRxCleanup();
|
|
|
|
/*!
|
|
* \brief Lower opaque constructs in TIRX programs: AllocBuffer, For(thread_binding),
|
|
* unit loop elimination. This is the tirx-specific counterpart of
|
|
* s_tir::LowerOpaqueBlock, without any SBlock handling.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerTIRxOpaque();
|
|
|
|
/*!
|
|
* \brief Lower the TIR to a lower level IR for the given target.
|
|
* \return The pass.
|
|
*/
|
|
TVM_DLL Pass LowerTIRx();
|
|
|
|
} // namespace transform
|
|
} // namespace tirx
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_TIR_TRANSFORM_H_
|