Merge pull request KhronosGroup#6325 from s-perron/id_overflow12

spirv-opt: Fix crashes in ConvertToHalfPass due to ID overflow

Author: s-perron

source/opt/amd_ext_to_khr.cpp

@@ -240,21 +240,25 @@ bool ReplaceSwizzleInvocations(IRContext* ctx, Instruction* inst,
   // This gives the offset in the group of 4 of this invocation.
   Instruction* quad_idx = ir_builder.AddBinaryOp(
       uint_type_id, spv::Op::OpBitwiseAnd, id->result_id(), quad_mask);
+  if (quad_idx == nullptr) return false;
 
   // Get the invocation id of the first invocation in the group of 4.
   Instruction* quad_ldr =
       ir_builder.AddBinaryOp(uint_type_id, spv::Op::OpBitwiseXor,
                              id->result_id(), quad_idx->result_id());
+  if (quad_ldr == nullptr) return false;
 
   // Get the offset of the target invocation from the offset vector.
   Instruction* my_offset =
       ir_builder.AddBinaryOp(uint_type_id, spv::Op::OpVectorExtractDynamic,
                              offset_id, quad_idx->result_id());
+  if (my_offset == nullptr) return false;
 
   // Determine the index of the invocation to read from.
   Instruction* target_inv =
       ir_builder.AddBinaryOp(uint_type_id, spv::Op::OpIAdd,
                              quad_ldr->result_id(), my_offset->result_id());
+  if (target_inv == nullptr) return false;
 
   // Do the group operations
   uint32_t uint_max_id = ir_builder.GetUintConstantId(0xFFFFFFFF);
@@ -368,13 +372,17 @@ bool ReplaceSwizzleInvocationsMasked(
   uint32_t mask_extended = ir_builder.GetUintConstantId(0xFFFFFFE0);
   Instruction* and_mask = ir_builder.AddBinaryOp(
       uint_type_id, spv::Op::OpBitwiseOr, uint_x, mask_extended);
+  if (and_mask == nullptr) return false;
   Instruction* and_result =
       ir_builder.AddBinaryOp(uint_type_id, spv::Op::OpBitwiseAnd,
                              id->result_id(), and_mask->result_id());
+  if (and_result == nullptr) return false;
   Instruction* or_result = ir_builder.AddBinaryOp(
       uint_type_id, spv::Op::OpBitwiseOr, and_result->result_id(), uint_y);
+  if (or_result == nullptr) return false;
   Instruction* target_inv = ir_builder.AddBinaryOp(
       uint_type_id, spv::Op::OpBitwiseXor, or_result->result_id(), uint_z);
+  if (target_inv == nullptr) return false;
 
   // Do the group operations
   uint32_t uint_max_id = ir_builder.GetUintConstantId(0xFFFFFFFF);
@@ -449,6 +457,7 @@ bool ReplaceWriteInvocation(IRContext* ctx, Instruction* inst,
   Instruction* cmp =
       ir_builder.AddBinaryOp(bool_type_id, spv::Op::OpIEqual, t->result_id(),
                              inst->GetSingleWordInOperand(4));
+  if (cmp == nullptr) return false;
 
   // Build a select.
   inst->SetOpcode(spv::Op::OpSelect);
@@ -524,6 +533,7 @@ bool ReplaceMbcnt(IRContext* context, Instruction* inst,
   Instruction* t =
       ir_builder.AddBinaryOp(mask_inst->type_id(), spv::Op::OpBitwiseAnd,
                              bitcast->result_id(), mask_id);
+  if (t == nullptr) return false;
 
   inst->SetOpcode(spv::Op::OpBitCount);
   inst->SetInOperands({{SPV_OPERAND_TYPE_ID, {t->result_id()}}});
@@ -631,10 +641,13 @@ bool ReplaceCubeFaceCoord(IRContext* ctx, Instruction* inst,
   // Find which values are negative.  Used in later computations.
   Instruction* is_z_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, z->result_id(), f0_const_id);
+  if (is_z_neg == nullptr) return false;
   Instruction* is_y_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, y->result_id(), f0_const_id);
+  if (is_y_neg == nullptr) return false;
   Instruction* is_x_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, x->result_id(), f0_const_id);
+  if (is_x_neg == nullptr) return false;
 
   // Compute cubema
   Instruction* amax_x_y = ir_builder.AddNaryExtendedInstruction(
@@ -645,19 +658,23 @@ bool ReplaceCubeFaceCoord(IRContext* ctx, Instruction* inst,
       {az->result_id(), amax_x_y->result_id()});
   Instruction* cubema = ir_builder.AddBinaryOp(float_type_id, spv::Op::OpFMul,
                                                f2_const_id, amax->result_id());
+  if (cubema == nullptr) return false;
 
   // Do the comparisons needed for computing cubesc and cubetc.
   Instruction* is_z_max =
       ir_builder.AddBinaryOp(bool_id, spv::Op::OpFOrdGreaterThanEqual,
                              az->result_id(), amax_x_y->result_id());
+  if (is_z_max == nullptr) return false;
   Instruction* not_is_z_max = ir_builder.AddUnaryOp(
       bool_id, spv::Op::OpLogicalNot, is_z_max->result_id());
   Instruction* y_gr_x =
       ir_builder.AddBinaryOp(bool_id, spv::Op::OpFOrdGreaterThanEqual,
                              ay->result_id(), ax->result_id());
+  if (y_gr_x == nullptr) return false;
   Instruction* is_y_max =
       ir_builder.AddBinaryOp(bool_id, spv::Op::OpLogicalAnd,
                              not_is_z_max->result_id(), y_gr_x->result_id());
+  if (is_y_max == nullptr) return false;
 
   // Select the correct value for cubesc.
   // TODO(1841): Handle id overflow.
@@ -691,6 +708,7 @@ bool ReplaceCubeFaceCoord(IRContext* ctx, Instruction* inst,
       v2_float_type_id, {cubema->result_id(), cubema->result_id()});
   Instruction* div = ir_builder.AddBinaryOp(
       v2_float_type_id, spv::Op::OpFDiv, cube->result_id(), denom->result_id());
+  if (div == nullptr) return false;
 
   // Get the final result by adding 0.5 to |div|.
   inst->SetOpcode(spv::Op::OpFAdd);
@@ -780,10 +798,13 @@ bool ReplaceCubeFaceIndex(IRContext* ctx, Instruction* inst,
   // Find which values are negative.  Used in later computations.
   Instruction* is_z_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, z->result_id(), f0_const_id);
+  if (is_z_neg == nullptr) return false;
   Instruction* is_y_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, y->result_id(), f0_const_id);
+  if (is_y_neg == nullptr) return false;
   Instruction* is_x_neg = ir_builder.AddBinaryOp(
       bool_id, spv::Op::OpFOrdLessThan, x->result_id(), f0_const_id);
+  if (is_x_neg == nullptr) return false;
 
   // Find the max value.
   Instruction* amax_x_y = ir_builder.AddNaryExtendedInstruction(
@@ -792,9 +813,11 @@ bool ReplaceCubeFaceIndex(IRContext* ctx, Instruction* inst,
   Instruction* is_z_max =
       ir_builder.AddBinaryOp(bool_id, spv::Op::OpFOrdGreaterThanEqual,
                              az->result_id(), amax_x_y->result_id());
+  if (is_z_max == nullptr) return false;
   Instruction* y_gr_x =
       ir_builder.AddBinaryOp(bool_id, spv::Op::OpFOrdGreaterThanEqual,
                              ay->result_id(), ax->result_id());
+  if (y_gr_x == nullptr) return false;
 
   // Get the value for each case.
   // TODO(1841): Handle id overflow.

source/opt/convert_to_half_pass.cpp

@@ -75,6 +75,9 @@ analysis::Type* ConvertToHalfPass::FloatScalarType(uint32_t width) {
 analysis::Type* ConvertToHalfPass::FloatVectorType(uint32_t v_len,
                                                    uint32_t width) {
   analysis::Type* reg_float_ty = FloatScalarType(width);
+  if (reg_float_ty == nullptr) {
+    return nullptr;
+  }
   analysis::Vector vec_ty(reg_float_ty, v_len);
   return context()->get_type_mgr()->GetRegisteredType(&vec_ty);
 }
@@ -85,6 +88,9 @@ analysis::Type* ConvertToHalfPass::FloatMatrixType(uint32_t v_cnt,
   Instruction* vty_inst = get_def_use_mgr()->GetDef(vty_id);
   uint32_t v_len = vty_inst->GetSingleWordInOperand(1);
   analysis::Type* reg_vec_ty = FloatVectorType(v_len, width);
+  if (reg_vec_ty == nullptr) {
+    return nullptr;
+  }
   analysis::Matrix mat_ty(reg_vec_ty, v_cnt);
   return context()->get_type_mgr()->GetRegisteredType(&mat_ty);
 }
@@ -99,6 +105,9 @@ uint32_t ConvertToHalfPass::EquivFloatTypeId(uint32_t ty_id, uint32_t width) {
     reg_equiv_ty = FloatVectorType(ty_inst->GetSingleWordInOperand(1), width);
   else  // spv::Op::OpTypeFloat
     reg_equiv_ty = FloatScalarType(width);
+  if (reg_equiv_ty == nullptr) {
+    return 0;
+  }
   return context()->get_type_mgr()->GetTypeInstruction(reg_equiv_ty);
 }
 
@@ -107,6 +116,10 @@ void ConvertToHalfPass::GenConvert(uint32_t* val_idp, uint32_t width,
   Instruction* val_inst = get_def_use_mgr()->GetDef(*val_idp);
   uint32_t ty_id = val_inst->type_id();
   uint32_t nty_id = EquivFloatTypeId(ty_id, width);
+  if (nty_id == 0) {
+    status_ = Status::Failure;
+    return;
+  }
   if (nty_id == ty_id) return;
   Instruction* cvt_inst;
   InstructionBuilder builder(
@@ -116,6 +129,10 @@ void ConvertToHalfPass::GenConvert(uint32_t* val_idp, uint32_t width,
     cvt_inst = builder.AddNullaryOp(nty_id, spv::Op::OpUndef);
   else
     cvt_inst = builder.AddUnaryOp(nty_id, spv::Op::OpFConvert, *val_idp);
+  if (cvt_inst == nullptr) {
+    status_ = Status::Failure;
+    return;
+  }
   *val_idp = cvt_inst->result_id();
 }
 
@@ -137,22 +154,43 @@ bool ConvertToHalfPass::MatConvertCleanup(Instruction* inst) {
   uint32_t orig_width = (cty_inst->GetSingleWordInOperand(0) == 16) ? 32 : 16;
   uint32_t orig_mat_id = inst->GetSingleWordInOperand(0);
   uint32_t orig_vty_id = EquivFloatTypeId(vty_id, orig_width);
+  if (orig_vty_id == 0) {
+    status_ = Status::Failure;
+    return false;
+  }
   std::vector<Operand> opnds = {};
   for (uint32_t vidx = 0; vidx < v_cnt; ++vidx) {
     Instruction* ext_inst = builder.AddIdLiteralOp(
         orig_vty_id, spv::Op::OpCompositeExtract, orig_mat_id, vidx);
+    if (ext_inst == nullptr) {
+      status_ = Status::Failure;
+      return false;
+    }
     Instruction* cvt_inst =
         builder.AddUnaryOp(vty_id, spv::Op::OpFConvert, ext_inst->result_id());
+    if (cvt_inst == nullptr) {
+      status_ = Status::Failure;
+      return false;
+    }
     opnds.push_back({SPV_OPERAND_TYPE_ID, {cvt_inst->result_id()}});
   }
   uint32_t mat_id = TakeNextId();
+  if (mat_id == 0) {
+    status_ = Status::Failure;
+    return false;
+  }
   std::unique_ptr<Instruction> mat_inst(new Instruction(
       context(), spv::Op::OpCompositeConstruct, mty_id, mat_id, opnds));
   (void)builder.AddInstruction(std::move(mat_inst));
   context()->ReplaceAllUsesWith(inst->result_id(), mat_id);
   // Turn original instruction into copy so it is valid.
+  uint32_t new_type_id = EquivFloatTypeId(mty_id, orig_width);
+  if (new_type_id == 0) {
+    status_ = Status::Failure;
+    return false;
+  }
   inst->SetOpcode(spv::Op::OpCopyObject);
-  inst->SetResultType(EquivFloatTypeId(mty_id, orig_width));
+  inst->SetResultType(new_type_id);
   get_def_use_mgr()->AnalyzeInstUse(inst);
   return true;
 }
@@ -187,13 +225,24 @@ bool ConvertToHalfPass::GenHalfArith(Instruction* inst) {
   // Convert all float32 based operands to float16 equivalent and change
   // instruction type to float16 equivalent.
   inst->ForEachInId([&inst, &modified, this](uint32_t* idp) {
+    if (status_ == Status::Failure) {
+      return;
+    }
     Instruction* op_inst = get_def_use_mgr()->GetDef(*idp);
     if (!IsFloat(op_inst, 32)) return;
     GenConvert(idp, 16, inst);
     modified = true;
   });
+  if (status_ == Status::Failure) {
+    return false;
+  }
   if (IsFloat(inst, 32)) {
-    inst->SetResultType(EquivFloatTypeId(inst->type_id(), 16));
+    uint32_t new_type_id = EquivFloatTypeId(inst->type_id(), 16);
+    if (new_type_id == 0) {
+      status_ = Status::Failure;
+      return false;
+    }
+    inst->SetResultType(new_type_id);
     converted_ids_.insert(inst->result_id());
     modified = true;
   }
@@ -211,6 +260,9 @@ bool ConvertToHalfPass::ProcessPhi(Instruction* inst, uint32_t from_width,
   bool modified = false;
   inst->ForEachInId([&ocnt, &prev_idp, &from_width, &to_width, &modified,
                      this](uint32_t* idp) {
+    if (status_ == Status::Failure) {
+      return;
+    }
     if (ocnt % 2 == 0) {
       prev_idp = idp;
     } else {
@@ -230,8 +282,16 @@ bool ConvertToHalfPass::ProcessPhi(Instruction* inst, uint32_t from_width,
     }
     ++ocnt;
   });
+  if (status_ == Status::Failure) {
+    return false;
+  }
   if (to_width == 16u) {
-    inst->SetResultType(EquivFloatTypeId(inst->type_id(), 16u));
+    uint32_t new_type_id = EquivFloatTypeId(inst->type_id(), 16u);
+    if (new_type_id == 0) {
+      status_ = Status::Failure;
+      return false;
+    }
+    inst->SetResultType(new_type_id);
     converted_ids_.insert(inst->result_id());
     modified = true;
   }
@@ -242,7 +302,12 @@ bool ConvertToHalfPass::ProcessPhi(Instruction* inst, uint32_t from_width,
 bool ConvertToHalfPass::ProcessConvert(Instruction* inst) {
   // If float32 and relaxed, change to float16 convert
   if (IsFloat(inst, 32) && IsRelaxed(inst->result_id())) {
-    inst->SetResultType(EquivFloatTypeId(inst->type_id(), 16));
+    uint32_t new_type_id = EquivFloatTypeId(inst->type_id(), 16);
+    if (new_type_id == 0) {
+      status_ = Status::Failure;
+      return false;
+    }
+    inst->SetResultType(new_type_id);
     get_def_use_mgr()->AnalyzeInstUse(inst);
     converted_ids_.insert(inst->result_id());
   }
@@ -255,7 +320,7 @@ bool ConvertToHalfPass::ProcessConvert(Instruction* inst) {
   Instruction* val_inst = get_def_use_mgr()->GetDef(val_id);
   if (inst->type_id() == val_inst->type_id())
     inst->SetOpcode(spv::Op::OpCopyObject);
-  return true;  // modified
+  return true;
 }
 
 bool ConvertToHalfPass::ProcessImageRef(Instruction* inst) {
@@ -265,6 +330,9 @@ bool ConvertToHalfPass::ProcessImageRef(Instruction* inst) {
     uint32_t dref_id = inst->GetSingleWordInOperand(kImageSampleDrefIdInIdx);
     if (converted_ids_.count(dref_id) > 0) {
       GenConvert(&dref_id, 32, inst);
+      if (status_ == Status::Failure) {
+        return false;
+      }
       inst->SetInOperand(kImageSampleDrefIdInIdx, {dref_id});
       get_def_use_mgr()->AnalyzeInstUse(inst);
       modified = true;
@@ -279,11 +347,17 @@ bool ConvertToHalfPass::ProcessDefault(Instruction* inst) {
   if (inst->opcode() == spv::Op::OpPhi) return ProcessPhi(inst, 16u, 32u);
   bool modified = false;
   inst->ForEachInId([&inst, &modified, this](uint32_t* idp) {
+    if (status_ == Status::Failure) {
+      return;
+    }
     if (converted_ids_.count(*idp) == 0) return;
     uint32_t old_id = *idp;
     GenConvert(idp, 32, inst);
     if (*idp != old_id) modified = true;
   });
+  if (status_ == Status::Failure) {
+    return false;
+  }
   if (modified) get_def_use_mgr()->AnalyzeInstUse(inst);
   return modified;
 }
@@ -370,19 +444,38 @@ bool ConvertToHalfPass::ProcessFunction(Function* func) {
       });
   // Replace invalid converts of matrix into equivalent vector extracts,
   // converts and finally a composite construct
+  bool ok = true;
   cfg()->ForEachBlockInReversePostOrder(
-      func->entry().get(), [&modified, this](BasicBlock* bb) {
-        for (auto ii = bb->begin(); ii != bb->end(); ++ii)
-          modified |= MatConvertCleanup(&*ii);
+      func->entry().get(), [&modified, &ok, this](BasicBlock* bb) {
+        if (!ok) {
+          return;
+        }
+        for (auto ii = bb->begin(); ii != bb->end(); ++ii) {
+          bool Mmodified = MatConvertCleanup(&*ii);
+          if (status_ == Status::Failure) {
+            ok = false;
+            break;
+          }
+          modified |= Mmodified;
+        }
       });
+
+  if (!ok) {
+    return false;
+  }
   return modified;
 }
 
 Pass::Status ConvertToHalfPass::ProcessImpl() {
+  status_ = Status::SuccessWithoutChange;
   Pass::ProcessFunction pfn = [this](Function* fp) {
     return ProcessFunction(fp);
   };
   bool modified = context()->ProcessReachableCallTree(pfn);
+  if (status_ == Status::Failure) {
+    return status_;
+  }
+
   // If modified, make sure module has Float16 capability
   if (modified) context()->AddCapability(spv::Capability::Float16);
   // Remove all RelaxedPrecision decorations from instructions and globals
@@ -514,4 +607,4 @@ void ConvertToHalfPass::Initialize() {
 }
 
 }  // namespace opt
-}  // namespace spvtools
+}  // namespace spvtools

source/opt/convert_to_half_pass.h

@@ -130,6 +130,9 @@ class ConvertToHalfPass : public Pass {
     }
   };
 
+  // The status of the pass.
+  Pass::Status status_;
+
   // Set of core operations to be processed
   std::unordered_set<spv::Op, hasher> target_ops_core_;