1// Copyright (c) 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#include "source/opt/reduce_load_size.h"
16
17#include <set>
18#include <vector>
19
20#include "source/opt/instruction.h"
21#include "source/opt/ir_builder.h"
22#include "source/opt/ir_context.h"
23#include "source/util/bit_vector.h"
24
25namespace {
26
27const uint32_t kExtractCompositeIdInIdx = 0;
28const uint32_t kVariableStorageClassInIdx = 0;
29const uint32_t kLoadPointerInIdx = 0;
30const double kThreshold = 0.9;
31
32} // namespace
33
34namespace spvtools {
35namespace opt {
36
37Pass::Status ReduceLoadSize::Process() {
38 bool modified = false;
39
40 for (auto& func : *get_module()) {
41 func.ForEachInst([&modified, this](Instruction* inst) {
42 if (inst->opcode() == SpvOpCompositeExtract) {
43 if (ShouldReplaceExtract(inst)) {
44 modified |= ReplaceExtract(inst);
45 }
46 }
47 });
48 }
49
50 return modified ? Status::SuccessWithChange : Status::SuccessWithoutChange;
51}
52
53bool ReduceLoadSize::ReplaceExtract(Instruction* inst) {
54 assert(inst->opcode() == SpvOpCompositeExtract &&
55 "Wrong opcode. Should be OpCompositeExtract.");
56 analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
57 analysis::TypeManager* type_mgr = context()->get_type_mgr();
58 analysis::ConstantManager* const_mgr = context()->get_constant_mgr();
59
60 uint32_t composite_id =
61 inst->GetSingleWordInOperand(kExtractCompositeIdInIdx);
62 Instruction* composite_inst = def_use_mgr->GetDef(composite_id);
63
64 if (composite_inst->opcode() != SpvOpLoad) {
65 return false;
66 }
67
68 analysis::Type* composite_type = type_mgr->GetType(composite_inst->type_id());
69 if (composite_type->kind() == analysis::Type::kVector ||
70 composite_type->kind() == analysis::Type::kMatrix) {
71 return false;
72 }
73
74 Instruction* var = composite_inst->GetBaseAddress();
75 if (var == nullptr || var->opcode() != SpvOpVariable) {
76 return false;
77 }
78
79 SpvStorageClass storage_class = static_cast<SpvStorageClass>(
80 var->GetSingleWordInOperand(kVariableStorageClassInIdx));
81 switch (storage_class) {
82 case SpvStorageClassUniform:
83 case SpvStorageClassUniformConstant:
84 case SpvStorageClassInput:
85 break;
86 default:
87 return false;
88 }
89
90 // Create a new access chain and load just after the old load.
91 // We cannot create the new access chain load in the position of the extract
92 // because the storage may have been written to in between.
93 InstructionBuilder ir_builder(
94 inst->context(), composite_inst,
95 IRContext::kAnalysisInstrToBlockMapping | IRContext::kAnalysisDefUse);
96
97 uint32_t pointer_to_result_type_id =
98 type_mgr->FindPointerToType(inst->type_id(), storage_class);
99 assert(pointer_to_result_type_id != 0 &&
100 "We did not find the pointer type that we need.");
101
102 analysis::Integer int_type(32, false);
103 const analysis::Type* uint32_type = type_mgr->GetRegisteredType(&int_type);
104 std::vector<uint32_t> ids;
105 for (uint32_t i = 1; i < inst->NumInOperands(); ++i) {
106 uint32_t index = inst->GetSingleWordInOperand(i);
107 const analysis::Constant* index_const =
108 const_mgr->GetConstant(uint32_type, {index});
109 ids.push_back(const_mgr->GetDefiningInstruction(index_const)->result_id());
110 }
111
112 Instruction* new_access_chain = ir_builder.AddAccessChain(
113 pointer_to_result_type_id,
114 composite_inst->GetSingleWordInOperand(kLoadPointerInIdx), ids);
115 Instruction* new_laod =
116 ir_builder.AddLoad(inst->type_id(), new_access_chain->result_id());
117
118 context()->ReplaceAllUsesWith(inst->result_id(), new_laod->result_id());
119 context()->KillInst(inst);
120 return true;
121}
122
123bool ReduceLoadSize::ShouldReplaceExtract(Instruction* inst) {
124 analysis::DefUseManager* def_use_mgr = context()->get_def_use_mgr();
125 Instruction* op_inst = def_use_mgr->GetDef(
126 inst->GetSingleWordInOperand(kExtractCompositeIdInIdx));
127
128 if (op_inst->opcode() != SpvOpLoad) {
129 return false;
130 }
131
132 auto cached_result = should_replace_cache_.find(op_inst->result_id());
133 if (cached_result != should_replace_cache_.end()) {
134 return cached_result->second;
135 }
136
137 bool all_elements_used = false;
138 std::set<uint32_t> elements_used;
139
140 all_elements_used =
141 !def_use_mgr->WhileEachUser(op_inst, [&elements_used](Instruction* use) {
142 if (use->opcode() != SpvOpCompositeExtract ||
143 use->NumInOperands() == 1) {
144 return false;
145 }
146 elements_used.insert(use->GetSingleWordInOperand(1));
147 return true;
148 });
149
150 bool should_replace = false;
151 if (all_elements_used) {
152 should_replace = false;
153 } else {
154 analysis::ConstantManager* const_mgr = context()->get_constant_mgr();
155 analysis::TypeManager* type_mgr = context()->get_type_mgr();
156 analysis::Type* load_type = type_mgr->GetType(op_inst->type_id());
157 uint32_t total_size = 1;
158 switch (load_type->kind()) {
159 case analysis::Type::kArray: {
160 const analysis::Constant* size_const =
161 const_mgr->FindDeclaredConstant(load_type->AsArray()->LengthId());
162 assert(size_const->AsIntConstant());
163 total_size = size_const->GetU32();
164 } break;
165 case analysis::Type::kStruct:
166 total_size = static_cast<uint32_t>(
167 load_type->AsStruct()->element_types().size());
168 break;
169 default:
170 break;
171 }
172 double percent_used = static_cast<double>(elements_used.size()) /
173 static_cast<double>(total_size);
174 should_replace = (percent_used < kThreshold);
175 }
176
177 should_replace_cache_[op_inst->result_id()] = should_replace;
178 return should_replace;
179}
180
181} // namespace opt
182} // namespace spvtools
183