Bindless Validation: Instrument descriptor-based loads and stores (#2583)

Essentially, support UBOs and SSBOs, scalar and array (sized and unsized).

source/opt/inst_bindless_check_pass.cpp

@@ -29,6 +29,7 @@ static const int kSpvAccessChainIndex0IdInIdx = 1;
 static const int kSpvTypePointerTypeIdInIdx = 1;
 static const int kSpvTypeArrayLengthIdInIdx = 1;
 static const int kSpvConstantValueInIdx = 0;
+static const int kSpvVariableStorageClassInIdx = 0;
 
 }  // anonymous namespace
 
@@ -59,36 +60,42 @@ uint32_t InstBindlessCheckPass::GenDebugReadInit(uint32_t var_id,
 
 uint32_t InstBindlessCheckPass::CloneOriginalReference(
     ref_analysis* ref, InstructionBuilder* builder) {
-  // Clone descriptor load
-  Instruction* load_inst = get_def_use_mgr()->GetDef(ref->load_id);
-  Instruction* new_load_inst =
-      builder->AddLoad(load_inst->type_id(),
-                       load_inst->GetSingleWordInOperand(kSpvLoadPtrIdInIdx));
-  uid2offset_[new_load_inst->unique_id()] = uid2offset_[load_inst->unique_id()];
-  uint32_t new_load_id = new_load_inst->result_id();
-  get_decoration_mgr()->CloneDecorations(load_inst->result_id(), new_load_id);
-  uint32_t new_image_id = new_load_id;
-  // Clone Image/SampledImage with new load, if needed
-  if (ref->image_id != 0) {
-    Instruction* image_inst = get_def_use_mgr()->GetDef(ref->image_id);
-    if (image_inst->opcode() == SpvOp::SpvOpSampledImage) {
-      Instruction* new_image_inst = builder->AddBinaryOp(
-          image_inst->type_id(), SpvOpSampledImage, new_load_id,
-          image_inst->GetSingleWordInOperand(kSpvSampledImageSamplerIdInIdx));
-      uid2offset_[new_image_inst->unique_id()] =
-          uid2offset_[image_inst->unique_id()];
-      new_image_id = new_image_inst->result_id();
-    } else {
-      assert(image_inst->opcode() == SpvOp::SpvOpImage && "expecting OpImage");
-      Instruction* new_image_inst =
-          builder->AddUnaryOp(image_inst->type_id(), SpvOpImage, new_load_id);
-      uid2offset_[new_image_inst->unique_id()] =
-          uid2offset_[image_inst->unique_id()];
-      new_image_id = new_image_inst->result_id();
+  // If original is image based, start by cloning descriptor load
+  uint32_t new_image_id = 0;
+  if (ref->desc_load_id != 0) {
+    Instruction* desc_load_inst = get_def_use_mgr()->GetDef(ref->desc_load_id);
+    Instruction* new_load_inst = builder->AddLoad(
+        desc_load_inst->type_id(),
+        desc_load_inst->GetSingleWordInOperand(kSpvLoadPtrIdInIdx));
+    uid2offset_[new_load_inst->unique_id()] =
+        uid2offset_[desc_load_inst->unique_id()];
+    uint32_t new_load_id = new_load_inst->result_id();
+    get_decoration_mgr()->CloneDecorations(desc_load_inst->result_id(),
+                                           new_load_id);
+    new_image_id = new_load_id;
+    // Clone Image/SampledImage with new load, if needed
+    if (ref->image_id != 0) {
+      Instruction* image_inst = get_def_use_mgr()->GetDef(ref->image_id);
+      if (image_inst->opcode() == SpvOp::SpvOpSampledImage) {
+        Instruction* new_image_inst = builder->AddBinaryOp(
+            image_inst->type_id(), SpvOpSampledImage, new_load_id,
+            image_inst->GetSingleWordInOperand(kSpvSampledImageSamplerIdInIdx));
+        uid2offset_[new_image_inst->unique_id()] =
+            uid2offset_[image_inst->unique_id()];
+        new_image_id = new_image_inst->result_id();
+      } else {
+        assert(image_inst->opcode() == SpvOp::SpvOpImage &&
+               "expecting OpImage");
+        Instruction* new_image_inst =
+            builder->AddUnaryOp(image_inst->type_id(), SpvOpImage, new_load_id);
+        uid2offset_[new_image_inst->unique_id()] =
+            uid2offset_[image_inst->unique_id()];
+        new_image_id = new_image_inst->result_id();
+      }
+      get_decoration_mgr()->CloneDecorations(ref->image_id, new_image_id);
     }
-    get_decoration_mgr()->CloneDecorations(ref->image_id, new_image_id);
   }
-  // Clone original reference using new image code
+  // Clone original reference
   std::unique_ptr<Instruction> new_ref_inst(ref->ref_inst->Clone(context()));
   uint32_t ref_result_id = ref->ref_inst->result_id();
   uint32_t new_ref_id = 0;
@@ -96,7 +103,9 @@ uint32_t InstBindlessCheckPass::CloneOriginalReference(
     new_ref_id = TakeNextId();
     new_ref_inst->SetResultId(new_ref_id);
   }
-  new_ref_inst->SetInOperand(kSpvImageSampleImageIdInIdx, {new_image_id});
+  // Update new ref with new image if created
+  if (new_image_id != 0)
+    new_ref_inst->SetInOperand(kSpvImageSampleImageIdInIdx, {new_image_id});
   // Register new reference and add to new block
   Instruction* added_inst = builder->AddInstruction(std::move(new_ref_inst));
   uid2offset_[added_inst->unique_id()] =
@@ -106,7 +115,7 @@ uint32_t InstBindlessCheckPass::CloneOriginalReference(
   return new_ref_id;
 }
 
-uint32_t InstBindlessCheckPass::GetDescriptorValueId(Instruction* inst) {
+uint32_t InstBindlessCheckPass::GetImageId(Instruction* inst) {
   switch (inst->opcode()) {
     case SpvOp::SpvOpImageSampleImplicitLod:
     case SpvOp::SpvOpImageSampleExplicitLod:
@@ -147,27 +156,77 @@ uint32_t InstBindlessCheckPass::GetDescriptorValueId(Instruction* inst) {
   return 0;
 }
 
+Instruction* InstBindlessCheckPass::GetDescriptorTypeInst(
+    Instruction* var_inst) {
+  uint32_t var_type_id = var_inst->type_id();
+  Instruction* var_type_inst = get_def_use_mgr()->GetDef(var_type_id);
+  uint32_t desc_type_id =
+      var_type_inst->GetSingleWordInOperand(kSpvTypePointerTypeIdInIdx);
+  return get_def_use_mgr()->GetDef(desc_type_id);
+}
+
 bool InstBindlessCheckPass::AnalyzeDescriptorReference(Instruction* ref_inst,
                                                        ref_analysis* ref) {
-  ref->image_id = GetDescriptorValueId(ref_inst);
+  ref->ref_inst = ref_inst;
+  if (ref_inst->opcode() == SpvOpLoad || ref_inst->opcode() == SpvOpStore) {
+    ref->desc_load_id = 0;
+    ref->ptr_id = ref_inst->GetSingleWordInOperand(kSpvLoadPtrIdInIdx);
+    Instruction* ptr_inst = get_def_use_mgr()->GetDef(ref->ptr_id);
+    if (ptr_inst->opcode() != SpvOp::SpvOpAccessChain) return false;
+    ref->var_id = ptr_inst->GetSingleWordInOperand(kSpvAccessChainBaseIdInIdx);
+    Instruction* var_inst = get_def_use_mgr()->GetDef(ref->var_id);
+    if (var_inst->opcode() != SpvOp::SpvOpVariable) return false;
+    uint32_t storage_class =
+        var_inst->GetSingleWordInOperand(kSpvVariableStorageClassInIdx);
+    switch (storage_class) {
+      case SpvStorageClassUniform:
+      case SpvStorageClassUniformConstant:
+      case SpvStorageClassStorageBuffer:
+        break;
+      default:
+        return false;
+        break;
+    }
+    Instruction* desc_type_inst = GetDescriptorTypeInst(var_inst);
+    switch (desc_type_inst->opcode()) {
+      case SpvOpTypeArray:
+      case SpvOpTypeRuntimeArray:
+        // A load through a descriptor array will have at least 3 operands. We
+        // do not want to instrument loads of descriptors here which are part of
+        // an image-based reference.
+        if (ptr_inst->NumInOperands() < 3) return false;
+        ref->index_id =
+            ptr_inst->GetSingleWordInOperand(kSpvAccessChainIndex0IdInIdx);
+        break;
+      default:
+        ref->index_id = 0;
+        break;
+    }
+    return true;
+  }
+  // Reference is not load or store. If not an image-based reference, return.
+  ref->image_id = GetImageId(ref_inst);
   if (ref->image_id == 0) return false;
   Instruction* image_inst = get_def_use_mgr()->GetDef(ref->image_id);
+  Instruction* desc_load_inst = nullptr;
   if (image_inst->opcode() == SpvOp::SpvOpSampledImage) {
-    ref->load_id =
+    ref->desc_load_id =
         image_inst->GetSingleWordInOperand(kSpvSampledImageImageIdInIdx);
+    desc_load_inst = get_def_use_mgr()->GetDef(ref->desc_load_id);
   } else if (image_inst->opcode() == SpvOp::SpvOpImage) {
-    ref->load_id =
+    ref->desc_load_id =
         image_inst->GetSingleWordInOperand(kSpvImageSampledImageIdInIdx);
+    desc_load_inst = get_def_use_mgr()->GetDef(ref->desc_load_id);
   } else {
-    ref->load_id = ref->image_id;
+    ref->desc_load_id = ref->image_id;
+    desc_load_inst = image_inst;
     ref->image_id = 0;
   }
-  Instruction* load_inst = get_def_use_mgr()->GetDef(ref->load_id);
-  if (load_inst->opcode() != SpvOp::SpvOpLoad) {
+  if (desc_load_inst->opcode() != SpvOp::SpvOpLoad) {
     // TODO(greg-lunarg): Handle additional possibilities?
     return false;
   }
-  ref->ptr_id = load_inst->GetSingleWordInOperand(kSpvLoadPtrIdInIdx);
+  ref->ptr_id = desc_load_inst->GetSingleWordInOperand(kSpvLoadPtrIdInIdx);
   Instruction* ptr_inst = get_def_use_mgr()->GetDef(ref->ptr_id);
   if (ptr_inst->opcode() == SpvOp::SpvOpVariable) {
     ref->index_id = 0;
@@ -189,7 +248,6 @@ bool InstBindlessCheckPass::AnalyzeDescriptorReference(Instruction* ref_inst,
     // TODO(greg-lunarg): Handle additional possibilities?
     return false;
   }
-  ref->ref_inst = ref_inst;
   return true;
 }
 
@@ -260,11 +318,7 @@ void InstBindlessCheckPass::GenBoundsCheckCode(
   // If index and bound both compile-time constants and index < bound,
   // return without changing
   Instruction* var_inst = get_def_use_mgr()->GetDef(ref.var_id);
-  uint32_t var_type_id = var_inst->type_id();
-  Instruction* var_type_inst = get_def_use_mgr()->GetDef(var_type_id);
-  uint32_t desc_type_id =
-      var_type_inst->GetSingleWordInOperand(kSpvTypePointerTypeIdInIdx);
-  Instruction* desc_type_inst = get_def_use_mgr()->GetDef(desc_type_id);
+  Instruction* desc_type_inst = GetDescriptorTypeInst(var_inst);
   uint32_t length_id = 0;
   if (desc_type_inst->opcode() == SpvOpTypeArray) {
     length_id =

source/opt/inst_bindless_check_pass.h

@@ -118,6 +118,7 @@ class InstBindlessCheckPass : public InstrumentPass {
   // AnalyzeDescriptorReference. It is necessary and sufficient for further
   // analysis and regeneration of the reference.
   typedef struct ref_analysis {
+    uint32_t desc_load_id;
     uint32_t image_id;
     uint32_t load_id;
     uint32_t ptr_id;
@@ -131,9 +132,12 @@ class InstBindlessCheckPass : public InstrumentPass {
   uint32_t CloneOriginalReference(ref_analysis* ref,
                                   InstructionBuilder* builder);
 
-  // If |inst| references through a descriptor, (ie references into an image
-  // or buffer), return the id of the value it references. Else return 0.
-  uint32_t GetDescriptorValueId(Instruction* inst);
+  // If |inst| references through an image, return the id of the image it
+  // references through. Else return 0.
+  uint32_t GetImageId(Instruction* inst);
+
+  // Get descriptor type inst of variable |var_inst|.
+  Instruction* GetDescriptorTypeInst(Instruction* var_inst);
 
   // Analyze descriptor reference |ref_inst| and save components into |ref|.
   // Return true if |ref_inst| is a descriptor reference, false otherwise.