spirv-opt: Handle id overflow in MergeReturnPass (#6340)
This CL adds error handling to the MergeReturnPass to gracefully handle
cases where the pass runs out of IDs.
The following functions were modified to return a boolean indicating
success or failure:
- AddNewPhiNodes
- AddReturnFlag
- AddReturnValue
- BranchToBlock
- CreatePhiNodesForInst
- ProcessStructuredBlock
- RecordReturned
- UpdatePhiNodes
The callers of these functions were updated to check the return value
and propagate the failure. This prevents the pass from crashing when it
runs out of IDs.
diff --git a/source/opt/merge_return_pass.cpp b/source/opt/merge_return_pass.cpp
index 427b4ec..47b0763 100644
--- a/source/opt/merge_return_pass.cpp
+++ b/source/opt/merge_return_pass.cpp
@@ -134,7 +134,9 @@
state_.pop_back();
}
- ProcessStructuredBlock(block);
+ if (!ProcessStructuredBlock(block)) {
+ return false;
+ }
// Generate state for next block if warranted
GenerateState(block);
@@ -169,7 +171,9 @@
// We have not kept the dominator tree up-to-date.
// Invalidate it at this point to make sure it will be rebuilt.
context()->RemoveDominatorAnalysis(function);
- AddNewPhiNodes();
+ if (!AddNewPhiNodes()) {
+ return false;
+ }
return true;
}
@@ -196,7 +200,9 @@
}
bool MergeReturnPass::CreateReturn(BasicBlock* block) {
- AddReturnValue();
+ if (!AddReturnValue()) {
+ return false;
+ }
if (return_value_) {
// Load and return the final return value
@@ -229,12 +235,18 @@
return true;
}
-void MergeReturnPass::ProcessStructuredBlock(BasicBlock* block) {
+bool MergeReturnPass::ProcessStructuredBlock(BasicBlock* block) {
+ if (block->tail() == block->end()) {
+ return true;
+ }
+
spv::Op tail_opcode = block->tail()->opcode();
if (tail_opcode == spv::Op::OpReturn ||
tail_opcode == spv::Op::OpReturnValue) {
if (!return_flag_) {
- AddReturnFlag();
+ if (!AddReturnFlag()) {
+ return false;
+ }
}
}
@@ -243,15 +255,20 @@
tail_opcode == spv::Op::OpUnreachable) {
assert(CurrentState().InBreakable() &&
"Should be in the placeholder construct.");
- BranchToBlock(block, CurrentState().BreakMergeId());
+ if (!BranchToBlock(block, CurrentState().BreakMergeId())) {
+ return false;
+ }
return_blocks_.insert(block->id());
}
+ return true;
}
-void MergeReturnPass::BranchToBlock(BasicBlock* block, uint32_t target) {
+bool MergeReturnPass::BranchToBlock(BasicBlock* block, uint32_t target) {
if (block->tail()->opcode() == spv::Op::OpReturn ||
block->tail()->opcode() == spv::Op::OpReturnValue) {
- RecordReturned(block);
+ if (!RecordReturned(block)) {
+ return false;
+ }
RecordReturnValue(block);
}
@@ -259,7 +276,9 @@
if (target_block->GetLoopMergeInst()) {
cfg()->SplitLoopHeader(target_block);
}
- UpdatePhiNodes(block, target_block);
+ if (!UpdatePhiNodes(block, target_block)) {
+ return false;
+ }
Instruction* return_inst = block->terminator();
return_inst->SetOpcode(spv::Op::OpBranch);
@@ -267,19 +286,26 @@
context()->get_def_use_mgr()->AnalyzeInstDefUse(return_inst);
new_edges_[target_block].insert(block->id());
cfg()->AddEdge(block->id(), target);
+ return true;
}
-void MergeReturnPass::UpdatePhiNodes(BasicBlock* new_source,
+bool MergeReturnPass::UpdatePhiNodes(BasicBlock* new_source,
BasicBlock* target) {
- target->ForEachPhiInst([this, new_source](Instruction* inst) {
+ bool succeeded = true;
+ target->ForEachPhiInst([this, new_source, &succeeded](Instruction* inst) {
uint32_t undefId = Type2Undef(inst->type_id());
+ if (undefId == 0) {
+ succeeded = false;
+ return;
+ }
inst->AddOperand({SPV_OPERAND_TYPE_ID, {undefId}});
inst->AddOperand({SPV_OPERAND_TYPE_ID, {new_source->id()}});
context()->UpdateDefUse(inst);
});
+ return succeeded;
}
-void MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
+bool MergeReturnPass::CreatePhiNodesForInst(BasicBlock* merge_block,
Instruction& inst) {
DominatorAnalysis* dom_tree =
context()->GetDominatorAnalysis(merge_block->GetParent());
@@ -313,7 +339,7 @@
});
if (users_to_update.empty()) {
- return;
+ return true;
}
// There is at least one values that needs to be replaced.
@@ -357,6 +383,9 @@
if (regenerateInstruction) {
std::unique_ptr<Instruction> regen_inst(inst.Clone(context()));
uint32_t new_id = TakeNextId();
+ if (new_id == 0) {
+ return false;
+ }
regen_inst->SetResultId(new_id);
Instruction* insert_pos = &*merge_block->begin();
while (insert_pos->opcode() == spv::Op::OpPhi) {
@@ -366,19 +395,31 @@
get_def_use_mgr()->AnalyzeInstDefUse(new_phi);
context()->set_instr_block(new_phi, merge_block);
- new_phi->ForEachInId([dom_tree, merge_block, this](uint32_t* use_id) {
+ bool succeeded = true;
+ new_phi->ForEachInId([dom_tree, merge_block, this,
+ &succeeded](uint32_t* use_id) {
+ if (!succeeded) {
+ return;
+ }
Instruction* use = get_def_use_mgr()->GetDef(*use_id);
BasicBlock* use_bb = context()->get_instr_block(use);
if (use_bb != nullptr && !dom_tree->Dominates(use_bb, merge_block)) {
- CreatePhiNodesForInst(merge_block, *use);
+ if (!CreatePhiNodesForInst(merge_block, *use)) {
+ succeeded = false;
+ }
}
});
+ if (!succeeded) {
+ return false;
+ }
} else {
InstructionBuilder builder(
context(), &*merge_block->begin(),
IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
- // TODO(1841): Handle id overflow.
new_phi = builder.AddPhi(inst.type_id(), phi_operands);
+ if (new_phi == nullptr) {
+ return false;
+ }
}
uint32_t result_of_phi = new_phi->result_id();
@@ -392,6 +433,7 @@
context()->AnalyzeUses(user);
}
}
+ return true;
}
bool MergeReturnPass::PredicateBlocks(
@@ -484,6 +526,9 @@
cfg()->RemoveSuccessorEdges(block);
auto old_body_id = TakeNextId();
+ if (old_body_id == 0) {
+ return false;
+ }
BasicBlock* old_body = block->SplitBasicBlock(context(), old_body_id, iter);
predicated->insert(old_body);
@@ -520,9 +565,11 @@
analysis::Bool bool_type;
uint32_t bool_id = context()->get_type_mgr()->GetId(&bool_type);
assert(bool_id != 0);
- // TODO(1841): Handle id overflow.
- uint32_t load_id =
- builder.AddLoad(bool_id, return_flag_->result_id())->result_id();
+ Instruction* load_inst = builder.AddLoad(bool_id, return_flag_->result_id());
+ if (load_inst == nullptr) {
+ return false;
+ }
+ uint32_t load_id = load_inst->result_id();
// 2. Branch to |merge_block| (true) or |old_body| (false)
builder.AddConditionalBranch(load_id, merge_block->id(), old_body->id(),
@@ -535,7 +582,9 @@
}
// 3. Update OpPhi instructions in |merge_block|.
- UpdatePhiNodes(block, merge_block);
+ if (!UpdatePhiNodes(block, merge_block)) {
+ return false;
+ }
// 4. Update the CFG. We do this after updating the OpPhi instructions
// because |UpdatePhiNodes| assumes the edge from |block| has not been added
@@ -548,10 +597,10 @@
return true;
}
-void MergeReturnPass::RecordReturned(BasicBlock* block) {
+bool MergeReturnPass::RecordReturned(BasicBlock* block) {
if (block->tail()->opcode() != spv::Op::OpReturn &&
block->tail()->opcode() != spv::Op::OpReturnValue)
- return;
+ return true;
assert(return_flag_ && "Did not generate the return flag variable.");
@@ -564,6 +613,9 @@
const analysis::Constant* true_const =
const_mgr->GetConstant(bool_type, {true});
constant_true_ = const_mgr->GetDefiningInstruction(true_const);
+ if (!constant_true_) {
+ return false;
+ }
context()->UpdateDefUse(constant_true_);
}
@@ -577,6 +629,7 @@
&*block->tail().InsertBefore(std::move(return_store));
context()->set_instr_block(store_inst, block);
context()->AnalyzeDefUse(store_inst);
+ return true;
}
void MergeReturnPass::RecordReturnValue(BasicBlock* block) {
@@ -600,18 +653,21 @@
context()->AnalyzeDefUse(store_inst);
}
-void MergeReturnPass::AddReturnValue() {
- if (return_value_) return;
+bool MergeReturnPass::AddReturnValue() {
+ if (return_value_) return true;
uint32_t return_type_id = function_->type_id();
if (get_def_use_mgr()->GetDef(return_type_id)->opcode() ==
spv::Op::OpTypeVoid)
- return;
+ return true;
uint32_t return_ptr_type = context()->get_type_mgr()->FindPointerToType(
return_type_id, spv::StorageClass::Function);
uint32_t var_id = TakeNextId();
+ if (var_id == 0) {
+ return false;
+ }
std::unique_ptr<Instruction> returnValue(
new Instruction(context(), spv::Op::OpVariable, return_ptr_type, var_id,
std::initializer_list<Operand>{
@@ -627,27 +683,44 @@
context()->get_decoration_mgr()->CloneDecorations(
function_->result_id(), var_id, {spv::Decoration::RelaxedPrecision});
+ return true;
}
-void MergeReturnPass::AddReturnFlag() {
- if (return_flag_) return;
+bool MergeReturnPass::AddReturnFlag() {
+ if (return_flag_) return true;
analysis::TypeManager* type_mgr = context()->get_type_mgr();
analysis::ConstantManager* const_mgr = context()->get_constant_mgr();
analysis::Bool temp;
uint32_t bool_id = type_mgr->GetTypeInstruction(&temp);
+ if (bool_id == 0) {
+ return false;
+ }
analysis::Bool* bool_type = type_mgr->GetType(bool_id)->AsBool();
const analysis::Constant* false_const =
const_mgr->GetConstant(bool_type, {false});
- uint32_t const_false_id =
- const_mgr->GetDefiningInstruction(false_const)->result_id();
+ Instruction* false_inst = const_mgr->GetDefiningInstruction(false_const);
+ if (false_inst == nullptr) {
+ return false;
+ }
+ uint32_t const_false_id = false_inst->result_id();
uint32_t bool_ptr_id =
type_mgr->FindPointerToType(bool_id, spv::StorageClass::Function);
+ if (bool_ptr_id == 0) {
+ return false;
+ ;
+ }
+
uint32_t var_id = TakeNextId();
+
+ if (var_id == 0) {
+ return false;
+ }
+
std::unique_ptr<Instruction> returnFlag(new Instruction(
context(), spv::Op::OpVariable, bool_ptr_id, var_id,
std::initializer_list<Operand>{{SPV_OPERAND_TYPE_STORAGE_CLASS,
@@ -661,6 +734,7 @@
return_flag_ = &*entry_block->begin();
context()->AnalyzeDefUse(return_flag_);
context()->set_instr_block(return_flag_, entry_block);
+ return true;
}
std::vector<BasicBlock*> MergeReturnPass::CollectReturnBlocks(
@@ -739,16 +813,19 @@
return true;
}
-void MergeReturnPass::AddNewPhiNodes() {
+bool MergeReturnPass::AddNewPhiNodes() {
std::list<BasicBlock*> order;
cfg()->ComputeStructuredOrder(function_, &*function_->begin(), &order);
for (BasicBlock* bb : order) {
- AddNewPhiNodes(bb);
+ if (!AddNewPhiNodes(bb)) {
+ return false;
+ }
}
+ return true;
}
-void MergeReturnPass::AddNewPhiNodes(BasicBlock* bb) {
+bool MergeReturnPass::AddNewPhiNodes(BasicBlock* bb) {
// New phi nodes are needed for any id whose definition used to dominate |bb|,
// but no longer dominates |bb|. These are found by walking the dominator
// tree starting at the original immediate dominator of |bb| and ending at its
@@ -766,16 +843,19 @@
BasicBlock* dominator = dom_tree->ImmediateDominator(bb);
if (dominator == nullptr) {
- return;
+ return true;
}
BasicBlock* current_bb = context()->get_instr_block(original_dominator_[bb]);
while (current_bb != nullptr && current_bb != dominator) {
for (Instruction& inst : *current_bb) {
- CreatePhiNodesForInst(bb, inst);
+ if (!CreatePhiNodesForInst(bb, inst)) {
+ return false;
+ }
}
current_bb = dom_tree->ImmediateDominator(current_bb);
}
+ return true;
}
void MergeReturnPass::RecordImmediateDominators(Function* function) {
@@ -859,8 +939,12 @@
++split_pos;
}
+ uint32_t new_block_id = TakeNextId();
+ if (new_block_id == 0) {
+ return false;
+ }
BasicBlock* old_block =
- start_block->SplitBasicBlock(context(), TakeNextId(), split_pos);
+ start_block->SplitBasicBlock(context(), new_block_id, split_pos);
// Find DebugFunctionDefinition inst in the old block, and if we can find it,
// move it to the entry block. Since DebugFunctionDefinition is not necessary
diff --git a/source/opt/merge_return_pass.h b/source/opt/merge_return_pass.h
index d39c1d9..d83ffc3 100644
--- a/source/opt/merge_return_pass.h
+++ b/source/opt/merge_return_pass.h
@@ -173,21 +173,22 @@
//
// Note this will break the semantics. To fix this, PredicateBlock will have
// to be called on the merge block the branch targets.
- void ProcessStructuredBlock(BasicBlock* block);
+ bool ProcessStructuredBlock(BasicBlock* block);
// Creates a variable used to store whether or not the control flow has
// traversed a block that used to have a return. A pointer to the instruction
- // declaring the variable is stored in |return_flag_|.
- void AddReturnFlag();
+ // declaring the variable is stored in |return_flag_|. Returns true if it
+ // succeeds.
+ bool AddReturnFlag();
// Creates the variable used to store the return value when passing through
- // a block that use to contain an OpReturnValue.
- void AddReturnValue();
+ // a block that use to contain an OpReturnValue. Returns true if it succeeds.
+ bool AddReturnValue();
- // Adds a store that stores true to |return_flag_| immediately before the
- // terminator of |block|. It is assumed that |AddReturnFlag| has already been
- // called.
- void RecordReturned(BasicBlock* block);
+ // Records that |block| used to be a return. This is done by adding an
+ // instruction to store true to the |return_flag_|. Returns true if it
+ // succeeds.
+ bool RecordReturned(BasicBlock* block);
// Adds an instruction that stores the value being returned in the
// OpReturnValue in |block|. The value is stored to |return_value_|, and the
@@ -198,10 +199,10 @@
// have already been called to create the variable to store to.
void RecordReturnValue(BasicBlock* block);
- // Adds an unconditional branch in |block| that branches to |target|. It also
- // adds stores to |return_flag_| and |return_value_| as needed.
- // |AddReturnFlag| and |AddReturnValue| must have already been called.
- void BranchToBlock(BasicBlock* block, uint32_t target);
+ // Replaces the terminator of |block| with a branch to |target|. If the
+ // terminator was a return, it will first call RecordReturned and
+ // RecordReturnValue. Returns true if it succeeds.
+ bool BranchToBlock(BasicBlock* block, uint32_t target);
// For every basic block that is reachable from |return_block|, extra code is
// added to jump around any code that should not be executed because the
@@ -239,32 +240,28 @@
// return block at the end of the pass.
bool CreateReturnBlock();
- // Creates a Phi node in |merge_block| for the result of |inst|.
- // Any uses of the result of |inst| that are no longer
- // dominated by |inst|, are replaced with the result of the new |OpPhi|
- // instruction.
- void CreatePhiNodesForInst(BasicBlock* merge_block, Instruction& inst);
+ // For each use of |inst| that is no longer dominated by |inst|, a phi node
+ // is created in |merge_block|. The original use is replaced by the result
+ // of the phi node. Returns true if it succeeds.
+ bool CreatePhiNodesForInst(BasicBlock* merge_block, Instruction& inst);
- // Add new phi nodes for any id that no longer dominate all of it uses. A phi
- // node is added to a block |bb| for an id if the id is defined between the
- // original immediate dominator of |bb| and its new immediate dominator. It
- // is assumed that at this point there are no unreachable blocks in the
- // control flow graph.
- void AddNewPhiNodes();
+ // Adds new phi nodes as needed to the function. This is necessary because
+ // adding the predication code can change the dominator tree. Returns false
+ // if there is a failure.
+ bool AddNewPhiNodes();
- // Creates any new phi nodes that are needed in |bb|. |AddNewPhiNodes| must
- // have already been called on the original dominators of |bb|.
- void AddNewPhiNodes(BasicBlock* bb);
+ // Adds new phi nodes to |bb| as needed. This is necessary because adding
+ // the predication code can change the dominator tree. Returns false if
+ // there is a failure.
+ bool AddNewPhiNodes(BasicBlock* bb);
// Records the terminator of immediate dominator for every basic block in
// |function|.
void RecordImmediateDominators(Function* function);
- // Modifies existing OpPhi instruction in |target| block to account for the
- // new edge from |new_source|. The value for that edge will be an Undef.
- //
- // The CFG must not include the edge from |new_source| to |target| yet.
- void UpdatePhiNodes(BasicBlock* new_source, BasicBlock* target);
+ // For each OpPhi instruction in |target|, this function adds an operand for
+ // |new_source|. The value will be OpUndef. Returns true if it succeeds.
+ bool UpdatePhiNodes(BasicBlock* new_source, BasicBlock* target);
StructuredControlState& CurrentState() { return state_.back(); }
@@ -334,4 +331,4 @@
} // namespace opt
} // namespace spvtools
-#endif // SOURCE_OPT_MERGE_RETURN_PASS_H_
\ No newline at end of file
+#endif // SOURCE_OPT_MERGE_RETURN_PASS_H_