1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file tir/usmp/transform.h
22 * \brief The transform passes for TIR-based Unified Static Memory Planner
23 */
24
25#ifndef TVM_TIR_USMP_TRANSFORM_H_
26#define TVM_TIR_USMP_TRANSFORM_H_
27
28#include <tvm/tir/usmp/utils.h>
29
30namespace tvm {
31namespace tir {
32namespace usmp {
33namespace transform {
34
35using Pass = tvm::transform::Pass;
36
37/*!
38 * \brief Convert the analyzed PoolAllocation to offsets from pool variables
39 *
40 * This pass would convert the main function to accept pool variables as an input
41 * that get passed onto the operator PrimFuncs. Furthermore, the static allocations
42 * will be converted to offsets within the pool variable.
43 *
44 * \return the pass
45 */
46TVM_DLL Pass ConvertPoolAllocationsToOffsets(const Map<tir::Stmt, PoolAllocation>& pool_allocations,
47 Bool emit_tvmscript_printable = Bool(false));
48
49/*!
50 * \brief Assign PoolInfo objects to tir.allocate nodes depending on the PrimFunc's target
51 *
52 * This pass would assign default PoolInfo objects to allocate nodes that are not otherwise
53 * annotated, depending on pool info supplied for each target.
54 *
55 * \return the pass
56 */
57TVM_DLL Pass AssignPoolInfo();
58
59/*!
60 * \brief This pass creates Allocate nodes for I/O tensors
61 *
62 * If the user wants to place the I/O tensors in the workspace, this pass is required to be
63 * run. In doing so, it will create Allocate nodes for I/O tensors to be planned, and be removed
64 * from function arguments.
65 *
66 * \return the pass
67 */
68TVM_DLL Pass CreateAllocatesForIO();
69
70} // namespace transform
71} // namespace usmp
72} // namespace tir
73} // namespace tvm
74
75#endif // TVM_TIR_USMP_TRANSFORM_H_
76