spirv-opt: Propagate failures in loop transformations (KhronosGroup#6330)

This CL modifies several loop transformation passes and their utility
functions to propagate failures, primarily caused by ID overflows when
creating new instructions or basic blocks.

Many functions that previously had a void return type now return a bool
or a status enum to indicate success or failure. The callers are updated
to check these return values and abort the transformation if a failure
occurs.

This affects loop unrolling, peeling, fusion, unswitching, and LICM,
as well as core IR cloning and manipulation functions.

The fuzzer transformations for inlining and outlining also get assertions
to detect cloning failures.

Author: s-perron

source/fuzz/transformation_inline_function.cpp

@@ -182,6 +182,8 @@ void TransformationInlineFunction::Apply(
     }
 
     auto* cloned_block = block.Clone(ir_context);
+    // TODO: Handle the nullptr.
+    assert(cloned_block);
     cloned_block = caller_function->InsertBasicBlockBefore(
         std::unique_ptr<opt::BasicBlock>(cloned_block), successor_block);
     cloned_block->GetLabel()->SetResultId(result_id_map.at(cloned_block->id()));

source/fuzz/transformation_outline_function.cpp

@@ -819,6 +819,8 @@ void TransformationOutlineFunction::PopulateOutlinedFunction(
     // Clone the block so that it can be added to the new function.
     auto cloned_block =
         std::unique_ptr<opt::BasicBlock>(block_it->Clone(ir_context));
+    // TODO: Handle a nullptr.
+    assert(cloned_block);
 
     // If this is the region's exit block, then the cloned block is the outlined
     // region's exit block.

source/opt/basic_block.cpp

@@ -29,11 +29,19 @@ constexpr uint32_t kSelectionMergeMergeBlockIdInIdx = 0;
 }  // namespace
 
 BasicBlock* BasicBlock::Clone(IRContext* context) const {
-  BasicBlock* clone = new BasicBlock(
-      std::unique_ptr<Instruction>(GetLabelInst()->Clone(context)));
+  Instruction* label_clone = GetLabelInst()->Clone(context);
+  if (!label_clone) {
+    return nullptr;
+  }
+  BasicBlock* clone = new BasicBlock(std::unique_ptr<Instruction>(label_clone));
   for (const auto& inst : insts_) {
     // Use the incoming context
-    clone->AddInstruction(std::unique_ptr<Instruction>(inst.Clone(context)));
+    Instruction* inst_clone = inst.Clone(context);
+    if (!inst_clone) {
+      delete clone;
+      return nullptr;
+    }
+    clone->AddInstruction(std::unique_ptr<Instruction>(inst_clone));
   }
 
   if (context->AreAnalysesValid(

source/opt/cfg.cpp

@@ -254,12 +254,13 @@ BasicBlock* CFG::SplitLoopHeader(BasicBlock* bb) {
   }
 
   // Adjust the OpPhi instructions as needed.
-  bb->ForEachPhiInst([latch_block, bb, new_header, context](Instruction* phi) {
+  bool ok = bb->WhileEachPhiInst([latch_block, bb, new_header,
+                                  context](Instruction* phi) -> bool {
     std::vector<uint32_t> preheader_phi_ops;
     std::vector<Operand> header_phi_ops;
 
-    // Identify where the original inputs to original OpPhi belong: header or
-    // preheader.
+    // Identify where the original inputs to original OpPhi belong: header
+    // or preheader.
     for (uint32_t i = 0; i < phi->NumInOperands(); i += 2) {
       uint32_t def_id = phi->GetSingleWordInOperand(i);
       uint32_t branch_id = phi->GetSingleWordInOperand(i + 1);
@@ -272,22 +273,24 @@ BasicBlock* CFG::SplitLoopHeader(BasicBlock* bb) {
       }
     }
 
-    // Create a phi instruction if and only if the preheader_phi_ops has more
-    // than one pair.
+    // Create a phi instruction if and only if the preheader_phi_ops has
+    // more than one pair.
     if (preheader_phi_ops.size() > 2) {
       InstructionBuilder builder(
           context, &*bb->begin(),
           IRContext::kAnalysisDefUse | IRContext::kAnalysisInstrToBlockMapping);
 
-      // TODO(1841): Handle id overflow.
       Instruction* new_phi = builder.AddPhi(phi->type_id(), preheader_phi_ops);
+      if (!new_phi) {
+        return false;
+      }
 
       // Add the OpPhi to the header bb.
       header_phi_ops.push_back({SPV_OPERAND_TYPE_ID, {new_phi->result_id()}});
       header_phi_ops.push_back({SPV_OPERAND_TYPE_ID, {bb->id()}});
     } else {
-      // An OpPhi with a single entry is just a copy.  In this case use the same
-      // instruction in the new header.
+      // An OpPhi with a single entry is just a copy.  In this case use the
+      // same instruction in the new header.
       header_phi_ops.push_back({SPV_OPERAND_TYPE_ID, {preheader_phi_ops[0]}});
       header_phi_ops.push_back({SPV_OPERAND_TYPE_ID, {bb->id()}});
     }
@@ -298,8 +301,13 @@ BasicBlock* CFG::SplitLoopHeader(BasicBlock* bb) {
     new_header->begin()->InsertBefore(std::move(phi_owner));
     context->set_instr_block(phi, new_header);
     context->AnalyzeUses(phi);
+    return true;
   });
 
+  if (!ok) {
+    return nullptr;
+  }
+
   // Add a branch to the new header.
   InstructionBuilder branch_builder(
       context, bb,

source/opt/function.cpp

@@ -40,6 +40,10 @@ Function* Function::Clone(IRContext* ctx) const {
   clone->blocks_.reserve(blocks_.size());
   for (const auto& b : blocks_) {
     std::unique_ptr<BasicBlock> bb(b->Clone(ctx));
+    if (!bb) {
+      delete clone;
+      return nullptr;
+    }
     clone->AddBasicBlock(std::move(bb));
   }
 

source/opt/instruction.cpp

@@ -168,7 +168,13 @@ Instruction* Instruction::Clone(IRContext* c) const {
   clone->dbg_line_insts_ = dbg_line_insts_;
   for (auto& i : clone->dbg_line_insts_) {
     i.unique_id_ = c->TakeNextUniqueId();
-    if (i.IsDebugLineInst()) i.SetResultId(c->TakeNextId());
+    if (i.IsDebugLineInst()) {
+      uint32_t new_id = c->TakeNextId();
+      if (new_id == 0) {
+        return nullptr;
+      }
+      i.SetResultId(new_id);
+    }
   }
   clone->dbg_scope_ = dbg_scope_;
   return clone;
@@ -546,27 +552,37 @@ void Instruction::ClearDbgLineInsts() {
   clear_dbg_line_insts();
 }
 
-void Instruction::UpdateDebugInfoFrom(const Instruction* from,
+bool Instruction::UpdateDebugInfoFrom(const Instruction* from,
                                       const Instruction* line) {
-  if (from == nullptr) return;
+  if (from == nullptr) return true;
   ClearDbgLineInsts();
   const Instruction* fromLine = line != nullptr ? line : from;
-  if (!fromLine->dbg_line_insts().empty())
-    AddDebugLine(&fromLine->dbg_line_insts().back());
+  if (!fromLine->dbg_line_insts().empty()) {
+    if (!AddDebugLine(&fromLine->dbg_line_insts().back())) {
+      return false;
+    }
+  }
   SetDebugScope(from->GetDebugScope());
   if (!IsLineInst() &&
       context()->AreAnalysesValid(IRContext::kAnalysisDebugInfo)) {
     context()->get_debug_info_mgr()->AnalyzeDebugInst(this);
   }
+  return true;
 }
 
-void Instruction::AddDebugLine(const Instruction* inst) {
+bool Instruction::AddDebugLine(const Instruction* inst) {
   dbg_line_insts_.push_back(*inst);
   dbg_line_insts_.back().unique_id_ = context()->TakeNextUniqueId();
-  if (inst->IsDebugLineInst())
-    dbg_line_insts_.back().SetResultId(context_->TakeNextId());
+  if (inst->IsDebugLineInst()) {
+    uint32_t new_id = context()->TakeNextId();
+    if (new_id == 0) {
+      return false;
+    }
+    dbg_line_insts_.back().SetResultId(new_id);
+  }
   if (context()->AreAnalysesValid(IRContext::kAnalysisDefUse))
     context()->get_def_use_mgr()->AnalyzeInstDefUse(&dbg_line_insts_.back());
+  return true;
 }
 
 bool Instruction::IsDebugLineInst() const {

source/opt/instruction.h

@@ -318,7 +318,7 @@ class Instruction : public utils::IntrusiveNodeBase<Instruction> {
   inline void SetDebugScope(const DebugScope& scope);
   inline const DebugScope& GetDebugScope() const { return dbg_scope_; }
   // Add debug line inst. Renew result id if Debug[No]Line
-  void AddDebugLine(const Instruction* inst);
+  bool AddDebugLine(const Instruction* inst);
   // Updates DebugInlinedAt of DebugScope and OpLine.
   void UpdateDebugInlinedAt(uint32_t new_inlined_at);
   // Clear line-related debug instructions attached to this instruction
@@ -338,7 +338,7 @@ class Instruction : public utils::IntrusiveNodeBase<Instruction> {
   // Updates lexical scope of DebugScope and OpLine.
   void UpdateLexicalScope(uint32_t scope);
   // Updates OpLine and DebugScope based on the information of |from|.
-  void UpdateDebugInfoFrom(const Instruction* from,
+  bool UpdateDebugInfoFrom(const Instruction* from,
                            const Instruction* line = nullptr);
   // Remove the |index|-th operand
   void RemoveOperand(uint32_t index) {

source/opt/licm_pass.cpp

@@ -118,7 +118,6 @@ bool LICMPass::IsImmediatelyContainedInLoop(Loop* loop, Function* f,
 }
 
 bool LICMPass::HoistInstruction(Loop* loop, Instruction* inst) {
-  // TODO(1841): Handle failure to create pre-header.
   BasicBlock* pre_header_bb = loop->GetOrCreatePreHeaderBlock();
   if (!pre_header_bb) {
     return false;

source/opt/loop_descriptor.cpp

@@ -25,6 +25,7 @@
 #include "source/opt/dominator_tree.h"
 #include "source/opt/ir_context.h"
 #include "source/opt/iterator.h"
+#include "source/opt/pass.h"
 #include "source/opt/tree_iterator.h"
 #include "source/util/make_unique.h"
 
@@ -278,6 +279,9 @@ BasicBlock* Loop::GetOrCreatePreHeaderBlock() {
 
   CFG* cfg = context_->cfg();
   loop_header_ = cfg->SplitLoopHeader(loop_header_);
+  if (loop_header_ == nullptr) {
+    return nullptr;
+  }
   return loop_preheader_;
 }
 
@@ -920,18 +924,19 @@ Instruction* Loop::FindConditionVariable(
   return induction;
 }
 
-bool LoopDescriptor::CreatePreHeaderBlocksIfMissing() {
-  auto modified = false;
+LoopDescriptor::Status LoopDescriptor::CreatePreHeaderBlocksIfMissing() {
+  bool modified = false;
 
   for (auto& loop : *this) {
     if (!loop.GetPreHeaderBlock()) {
+      if (!loop.GetOrCreatePreHeaderBlock()) {
+        return Status::kFailure;
+      }
       modified = true;
-      // TODO(1841): Handle failure to create pre-header.
-      loop.GetOrCreatePreHeaderBlock();
     }
   }
 
-  return modified;
+  return modified ? Status::kSuccessWithChange : Status::kSuccessWithoutChange;
 }
 
 // Add and remove loops which have been marked for addition and removal to

source/opt/loop_descriptor.h

@@ -425,6 +425,9 @@ class LoopDescriptor {
   using pre_iterator = TreeDFIterator<Loop>;
   using const_pre_iterator = TreeDFIterator<const Loop>;
 
+  // The status of processing a module.
+  enum class Status { kSuccessWithChange, kSuccessWithoutChange, kFailure };
+
   // Creates a loop object for all loops found in |f|.
   LoopDescriptor(IRContext* context, const Function* f);
 
@@ -506,9 +509,11 @@ class LoopDescriptor {
     loops_to_add_.emplace_back(std::make_pair(parent, std::move(loop_to_add)));
   }
 
-  // Checks all loops in |this| and will create pre-headers for all loops
-  // that don't have one. Returns |true| if any blocks were created.
-  bool CreatePreHeaderBlocksIfMissing();
+  // Creates pre-header blocks for all loops in the function that do not have
+  // one. Returns `LoopDescriptor::Status::kSuccessWithChange` if any change is
+  // made, `LoopDescriptor::Status::kSuccessWithoutChange` if no change is made,
+  // and `LoopDescriptor::Status::kFailure` if it fails to create a pre-header.
+  Status CreatePreHeaderBlocksIfMissing();
 
   // Should be called to preserve the LoopAnalysis after loops have been marked
   // for addition with AddLoop or MarkLoopForRemoval.