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 side_effect.cc
22 * \brief side effect analysis
23 */
24#include <tvm/ir/op.h>
25#include <tvm/tir/analysis.h>
26#include <tvm/tir/expr.h>
27#include <tvm/tir/expr_functor.h>
28#include <tvm/tir/op_attr_types.h>
29
30namespace tvm {
31namespace tir {
32
33class ExprSideEffect : public ExprVisitor {
34 public:
35 void VisitExpr(const PrimExpr& e) final {
36 if (kind_ == CallEffectKind::kUpdateState) return;
37 ExprVisitor::VisitExpr(e);
38 }
39
40 void VisitExpr_(const LoadNode* op) final {
41 this->UpdateEffect(CallEffectKind::kReadState);
42 ExprVisitor::VisitExpr_(op);
43 }
44
45 void VisitExpr_(const BufferLoadNode* op) final {
46 this->UpdateEffect(CallEffectKind::kReadState);
47 ExprVisitor::VisitExpr_(op);
48 }
49
50 void VisitExpr_(const CallNode* op) final {
51 static auto op_call_effect = Op::GetAttrMap<TCallEffectKind>("TCallEffectKind");
52
53 if (auto* ptr_op = op->op.as<OpNode>()) {
54 this->UpdateEffect(static_cast<CallEffectKind>(op_call_effect[GetRef<Op>(ptr_op)]->value));
55 } else {
56 this->UpdateEffect(CallEffectKind::kOpaque);
57 }
58 ExprVisitor::VisitExpr_(op);
59 }
60
61 void UpdateEffect(CallEffectKind effect_kind) {
62 if (effect_kind > CallEffectKind::kUpdateState) {
63 effect_kind = CallEffectKind::kUpdateState;
64 }
65 if (effect_kind > kind_) {
66 kind_ = effect_kind;
67 }
68 }
69
70 CallEffectKind kind_{CallEffectKind::kPure};
71};
72
73CallEffectKind SideEffect(const PrimExpr& e) {
74 ExprSideEffect visitor;
75 visitor(e);
76 return visitor.kind_;
77}
78
79} // namespace tir
80} // namespace tvm
81