Fixed a crash with cbuffer lowering with new poison value mechanism (microsoft#4684)

There was a crash when CBuffer vector subscript uses a non-literal index but is replaced with Constant in the optimizer before CBuffer load is lowerd. The CBuffer lowering code expects the constant to be within bounds, since it's translating the GEP to extractvalue. The out of bounds subscript is not correct and should cause an error, but ONLY if it actually exists in the final DXIL.

This change adds a way to emit error that only get emitted if the associated code is not removed by the end of compilation. Instead of emitting an error right away, emit a poison value with line number and error message instead and have it used by the problematic code. If the problematic code is not removed by the end of compilation, then this poison value would also be there, and it's safe to emit a real error based on it.

Author: adam-yang

include/dxc/HLSL/DxilPoisonValues.h

@@ -0,0 +1,37 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// DxilPoisonValues.h                                                        //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Allows insertion of poisoned values with error messages that get          //
+// cleaned up late in the compiler.                                          //
+//                                                                           //
+/////////////////////////////////////////////////////////////////////////////// 
+#pragma once
+
+#include "llvm/IR/DebugLoc.h"
+
+namespace llvm {
+  class Instruction;
+  class Type;
+  class Value;
+  class Module;
+}
+
+namespace hlsl {
+  // Create a special dx.poison.* instruction with debug location and an error message.
+  // The reason for this is certain invalid code is allowed in the code as long as it is
+  // removed by optimization at the end of compilation. We only want to emit the error
+  // for real if we are sure the code with the problem is in the final DXIL.
+  // 
+  // This "emits" an error message with the specified type. If by the end the compilation
+  // it is still used, then FinalizePoisonValues will emit the correct error for real.
+  llvm::Value *CreatePoisonValue(llvm::Type *ty, const llvm::Twine &errMsg, llvm::DebugLoc DL, llvm::Instruction *InsertPt);
+
+  // If there's any dx.poison.* values present in the module, emit them. Returns true
+  // M was modified in any way.
+  bool FinalizePoisonValues(llvm::Module &M);
+}
+

lib/HLSL/CMakeLists.txt

@@ -58,6 +58,7 @@ add_llvm_library(LLVMHLSL
   PauseResumePasses.cpp
   WaveSensitivityAnalysis.cpp
   DxilNoOptLegalize.cpp
+  DxilPoisonValues.cpp
   DxilDeleteRedundantDebugValues.cpp
 
   ADDITIONAL_HEADER_DIRS

lib/HLSL/DxilPoisonValues.cpp

@@ -0,0 +1,75 @@
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// DxilPoisonValues.cpp                                                      //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+// Allows insertion of poisoned values with error messages that get          //
+// cleaned up late in the compiler.                                          //
+//                                                                           //
+/////////////////////////////////////////////////////////////////////////////// 
+
+#include "dxc/HLSL/DxilPoisonValues.h"
+
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/LLVMContext.h"
+
+using namespace llvm;
+
+constexpr const char kPoisonPrefix[] = "dx.poison.";
+
+namespace hlsl {
+  Value *CreatePoisonValue(Type *ty, const Twine &errMsg, DebugLoc DL, Instruction *InsertPt) {
+    std::string functionName;
+    {
+      llvm::raw_string_ostream os(functionName);
+      os << kPoisonPrefix;
+      os << *ty;
+      os.flush();
+    }
+
+    Module &M = *InsertPt->getModule();
+
+    LLVMContext &C = M.getContext();
+    Type *argTypes[] = { Type::getMetadataTy(C) };
+    FunctionType *ft = FunctionType::get(ty, argTypes, false);
+    Constant *f = M.getOrInsertFunction(functionName, ft);
+
+    std::string errMsgStr = errMsg.str();
+    Value *args[] = {
+      MetadataAsValue::get(C, MDString::get(C, errMsgStr))
+    };
+    CallInst *ret = CallInst::Create(f, ArrayRef<Value *>(args), "err", InsertPt);
+    ret->setDebugLoc(DL);
+    return ret;
+  }
+
+  bool FinalizePoisonValues(Module &M) {
+    bool changed = false;
+    LLVMContext &Ctx = M.getContext();
+    for (auto it = M.begin(); it != M.end();) {
+      Function *F = &*(it++);
+      if (F->getName().startswith(kPoisonPrefix)) {
+        for (auto it = F->user_begin(); it != F->user_end();) {
+          User *U = *(it++);
+          CallInst *call = cast<CallInst>(U);
+          MDString *errMsgMD = cast<MDString>(cast<MetadataAsValue>(call->getArgOperand(0))->getMetadata());
+          StringRef errMsg = errMsgMD->getString();
+
+          Ctx.diagnose(DiagnosticInfoDxil(F, call->getDebugLoc(), errMsg, DS_Error));
+          if (!call->getType()->isVoidTy())
+            call->replaceAllUsesWith(UndefValue::get(call->getType()));
+          call->eraseFromParent();
+        }
+        F->eraseFromParent();
+        changed = true;
+      }
+    }
+    return changed;
+  }
+}
+

lib/HLSL/DxilPreparePasses.cpp

@@ -21,6 +21,7 @@
 #include "dxc/DXIL/DxilInstructions.h"
 #include "dxc/DXIL/DxilConstants.h"
 #include "dxc/HlslIntrinsicOp.h"
+#include "dxc/HLSL/DxilPoisonValues.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
@@ -757,6 +758,10 @@ class DxilFinalizeModule : public ModulePass {
   }
 
   bool runOnModule(Module &M) override {
+
+    // Remove all the poisoned values and emit errors if necessary.
+    (void)hlsl::FinalizePoisonValues(M);
+
     if (M.HasDxilModule()) {
       DxilModule &DM = M.GetDxilModule();
       unsigned ValMajor = 0;

lib/HLSL/HLOperationLower.cpp

@@ -26,6 +26,7 @@
 #include "dxc/HLSL/HLOperations.h"
 #include "dxc/HlslIntrinsicOp.h"
 #include "dxc/DXIL/DxilResourceProperties.h"
+#include "dxc/HLSL/DxilPoisonValues.h"
 
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/IRBuilder.h"
@@ -135,6 +136,34 @@ struct HLObjectOperationLowerHelper {
     MarkHasCounterOnCreateHandle(counterHandle, resSet);
   }
 
+  DxilResourceBase *FindCBufferResourceFromHandle(Value *handle) {
+    if (CallInst *CI = dyn_cast<CallInst>(handle)) {
+      hlsl::HLOpcodeGroup group =
+          hlsl::GetHLOpcodeGroupByName(CI->getCalledFunction());
+      if (group == HLOpcodeGroup::HLAnnotateHandle) {
+        handle = CI->getArgOperand(HLOperandIndex::kAnnotateHandleHandleOpIdx);
+      }
+    }
+
+    Constant *symbol = nullptr;
+    if (CallInst *CI = dyn_cast<CallInst>(handle)) {
+      hlsl::HLOpcodeGroup group =
+          hlsl::GetHLOpcodeGroupByName(CI->getCalledFunction());
+      if (group == HLOpcodeGroup::HLCreateHandle) {
+        symbol = dyn_cast<Constant>(CI->getArgOperand(HLOperandIndex::kCreateHandleResourceOpIdx));
+      }
+    }
+
+    if (!symbol)
+      return nullptr;
+
+    for (const std::unique_ptr<DxilCBuffer> &res : HLM.GetCBuffers()) {
+      if (res->GetGlobalSymbol() == symbol)
+        return res.get();
+    }
+    return nullptr;
+  }
+
   Value *GetOrCreateResourceForCbPtr(GetElementPtrInst *CbPtr,
                                      GlobalVariable *CbGV,
                                      DxilResourceProperties &RP) {
@@ -6816,31 +6845,33 @@ void TranslateCBGepLegacy(GetElementPtrInst *GEP, Value *handle,
       // Array always start from x channel.
       channel = 0;
     } else if (GEPIt->isVectorTy()) {
-      unsigned size = DL.getTypeAllocSize(GEPIt->getVectorElementType());
       // Indexing on vector.
       if (bImmIdx) {
-        unsigned tempOffset = size * immIdx;
-        if (size == 2) { // 16-bit types
-          unsigned channelInc = tempOffset >> 1;
-          DXASSERT((channel + channelInc) <= 8, "vector should not cross cb register (8x16bit)");
+        if (immIdx < GEPIt->getVectorNumElements()) {
+          const unsigned vectorElmSize     = DL.getTypeAllocSize(GEPIt->getVectorElementType());
+          const bool     bIs16bitType      = vectorElmSize == 2;
+          const unsigned tempOffset        = vectorElmSize * immIdx;
+          const unsigned numChannelsPerRow = bIs16bitType ? 8 : 4;
+          const unsigned channelInc        = bIs16bitType ? tempOffset >> 1 : tempOffset >> 2;
+
+          DXASSERT((channel + channelInc) < numChannelsPerRow, "vector should not cross cb register");
           channel += channelInc;
-          if (channel == 8) {
+          if (channel == numChannelsPerRow) {
             // Get to another row.
             // Update index and channel.
             channel = 0;
             legacyIndex = Builder.CreateAdd(legacyIndex, Builder.getInt32(1));
           }
         }
         else {
-          unsigned channelInc = tempOffset >> 2;
-          DXASSERT((channel + channelInc) <= 4, "vector should not cross cb register (8x32bit)");
-          channel += channelInc;
-          if (channel == 4) {
-            // Get to another row.
-            // Update index and channel.
-            channel = 0;
-            legacyIndex = Builder.CreateAdd(legacyIndex, Builder.getInt32(1));
+          StringRef resName = "(unknown)";
+          if (DxilResourceBase *Res = pObjHelper->FindCBufferResourceFromHandle(handle)) {
+            resName = Res->GetGlobalName();
           }
+          legacyIndex = hlsl::CreatePoisonValue(legacyIndex->getType(),
+            Twine("Out of bounds index (") + Twine(immIdx) + Twine(") in CBuffer '") + Twine(resName) + ("'"),
+            GEP->getDebugLoc(), GEP);
+          channel = 0;
         }
       } else {
         Type *EltTy = GEPIt->getVectorElementType();

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/vector_subscript.hlsl

@@ -0,0 +1,25 @@
+// RUN: %dxc -E main -T ps_6_0 -Od %s | FileCheck %s
+
+cbuffer constants : register(b0)
+{
+    float4 foo;
+}
+
+float main() : SV_TARGET
+{
+    float ret = 0;
+    uint index = 5;
+    if (index < 4) {
+        ret += foo[index];
+    }
+    return ret;
+}
+
+// Regression test for Od when an out of bound access happens in a dead block.
+// We want to make sure it doesn't crash, compiles successfully, and the
+// cbuffer load doesn't actually get generated.
+
+// CHECK: @main() {
+// CHECK-NOT: @dx.op.cbufferLoadLegacy.f32(i32 59
+// CHECK: call void @dx.op.storeOutput.f32(
+

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/vector_subscript_16bit.hlsl

@@ -0,0 +1,24 @@
+// RUN: %dxc -E main -T ps_6_2 -enable-16bit-types -Od %s | FileCheck %s
+
+cbuffer constants : register(b0)
+{
+    half4 foo;
+}
+
+float main() : SV_TARGET
+{
+    float ret = 0;
+    uint index = 9;
+    if (index < 8) {
+        ret += foo[index];
+    }
+    return ret;
+}
+
+// Regression test for Od when an out of bound access happens in a dead block.
+// We want to make sure it doesn't crash, compiles successfully, and the
+// cbuffer load doesn't actually get generated.
+
+// CHECK: @main() {
+// CHECK-NOT: @dx.op.cbufferLoadLegacy.f32(i32 59
+// CHECK: call void @dx.op.storeOutput.f32(

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/vector_subscript_16bit_error.hlsl

@@ -0,0 +1,24 @@
+// RUN: %dxc -E main -T ps_6_2 -enable-16bit-types -Od %s | FileCheck %s
+
+cbuffer constants : register(b0)
+{
+    half2 foo;
+}
+
+float main() : SV_TARGET
+{
+    float ret = 0;
+    uint index = 2;
+    if (index < 10) {
+        ret += foo[index];
+    }
+    return ret;
+}
+
+// Regression test for Od when an out of bound access happens but:
+// 1) the front-end can't immediate figure it out,
+// 2) when lowering the cbuffer load, the index is resolved to be a constant
+// 3) lowering code crashed because it assumed OOB access isn't possible when the index is constant.
+
+// CHECK: 13:{{[0-9]+}}: error: Out of bounds index (2) in CBuffer 'constants'
+

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/vector_subscript_error.hlsl

@@ -0,0 +1,25 @@
+// RUN: %dxc -E main -T ps_6_0 -Od %s | FileCheck %s
+
+cbuffer constants : register(b0)
+{
+    float2 foo;
+}
+
+float main() : SV_TARGET
+{
+    float ret = 0;
+    uint index = 2;
+    if (index < 10) {
+        ret += foo[index];
+    }
+    return ret;
+}
+
+// Regression test for Od when an out of bound access happens but:
+// 1) the front-end can't immediate figure it out,
+// 2) when lowering the cbuffer load, the index is resolved to be a constant
+// 3) lowering code crashed because it assumed OOB access isn't possible when the index is constant.
+
+// CHECK: 13:{{[0-9]+}}: error: Out of bounds index (2) in CBuffer 'constants'
+
+

tools/clang/test/HLSLFileCheck/hlsl/objects/CbufferLegacy/vector_subscript_error_noscope.hlsl

@@ -0,0 +1,22 @@
+// RUN: %dxc -E main -T ps_6_0 -Od %s | FileCheck %s
+
+const float4 foo;
+
+float main() : SV_TARGET
+{
+    float ret = 0;
+    uint index = 5;
+    if (index < 10) {
+        ret += foo[index];
+    }
+    return ret;
+}
+
+// Regression test for Od when an out of bound access happens but:
+// 1) the front-end can't immediate figure it out,
+// 2) when lowering the cbuffer load, the index is resolved to be a constant
+// 3) lowering code crashed because it assumed OOB access isn't possible when the index is constant.
+
+// CHECK: 10:{{[0-9]+}}: error: Out of bounds index (5) in CBuffer '$Globals'
+
+