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 update_pointer_storage_scope.h
22 * \brief A pass to update storage scopes for buffer variables.
23 */
24#ifndef TVM_TIR_TRANSFORMS_UPDATE_POINTER_STORAGE_SCOPE_H_
25#define TVM_TIR_TRANSFORMS_UPDATE_POINTER_STORAGE_SCOPE_H_
26
27#include <tvm/tir/expr.h>
28#include <tvm/tir/op.h>
29#include <tvm/tir/stmt_functor.h>
30
31#include <unordered_map>
32
33namespace tvm {
34namespace tir {
35
36class UpdatePointerStorageScope : public StmtExprMutator {
37 public:
38 explicit UpdatePointerStorageScope(
39 const std::unordered_map<const VarNode*, String>& new_storage_scopes);
40
41 virtual PrimExpr VisitExpr_(const VarNode*);
42 virtual PrimExpr VisitExpr_(const LoadNode*);
43 virtual PrimExpr VisitExpr_(const BufferLoadNode*);
44 virtual Stmt VisitStmt_(const AllocateNode*);
45 virtual Stmt VisitStmt_(const StoreNode*);
46 virtual Stmt VisitStmt_(const BufferStoreNode*);
47
48 private:
49 template <typename Node>
50 Node UpdateBufferAccess(Node node);
51
52 Buffer GetUpdatedBuffer(Buffer buf);
53
54 std::unordered_map<const VarNode*, Var> new_var_remap_;
55 std::unordered_map<const BufferNode*, Buffer> new_buffer_remap_;
56};
57
58} // namespace tir
59} // namespace tvm
60#endif // TVM_TIR_TRANSFORMS_UPDATE_POINTER_STORAGE_SCOPE_H_
61