[DirectX] Fix the writing of ConstantExpr GEPs to DXIL bitcode

llvm/lib/Target/DirectX/DXILDataScalarization.cpp

@@ -153,6 +153,8 @@ bool DataScalarizerVisitor::visitLoadInst(LoadInst &LI) {
     NewLoad->setAlignment(LI.getAlign());
     LI.replaceAllUsesWith(NewLoad);
     LI.eraseFromParent();
+    if (CE->use_empty())
+      CE->destroyConstant();
     visitGetElementPtrInst(*OldGEP);
     return true;
   }
@@ -173,6 +175,8 @@ bool DataScalarizerVisitor::visitStoreInst(StoreInst &SI) {
     NewStore->setAlignment(SI.getAlign());
     SI.replaceAllUsesWith(NewStore);
     SI.eraseFromParent();
+    if (CE->use_empty())
+      CE->destroyConstant();
     visitGetElementPtrInst(*OldGEP);
     return true;
   }
@@ -343,7 +347,7 @@ bool DataScalarizerVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
 
   GOp->replaceAllUsesWith(NewGEP);
 
-  if (auto *CE = dyn_cast<ConstantExpr>(GOp))
+  if (auto *CE = dyn_cast<ConstantExpr>(GOp); CE && CE->use_empty())
     CE->destroyConstant();
   else if (auto *OldGEPI = dyn_cast<GetElementPtrInst>(GOp))
     OldGEPI->eraseFromParent();

llvm/lib/Target/DirectX/DXILFlattenArrays.cpp

@@ -166,6 +166,8 @@ bool DXILFlattenArraysVisitor::visitLoadInst(LoadInst &LI) {
       NewLoad->setAlignment(LI.getAlign());
       LI.replaceAllUsesWith(NewLoad);
       LI.eraseFromParent();
+      if (CE->use_empty())
+        CE->destroyConstant();
       visitGetElementPtrInst(*OldGEP);
       return true;
     }
@@ -188,6 +190,8 @@ bool DXILFlattenArraysVisitor::visitStoreInst(StoreInst &SI) {
       NewStore->setAlignment(SI.getAlign());
       SI.replaceAllUsesWith(NewStore);
       SI.eraseFromParent();
+      if (CE->use_empty())
+        CE->destroyConstant();
       visitGetElementPtrInst(*OldGEP);
       return true;
     }
@@ -243,6 +247,8 @@ bool DXILFlattenArraysVisitor::visitGetElementPtrInst(GetElementPtrInst &GEP) {
 
     GEP.replaceAllUsesWith(NewGEPI);
     GEP.eraseFromParent();
+    if (PtrOpGEPCE->use_empty())
+      PtrOpGEPCE->destroyConstant();
     visitGetElementPtrInst(*OldGEPI);
     visitGetElementPtrInst(*NewGEPI);
     return true;

llvm/lib/Target/DirectX/DXILLegalizePass.cpp

@@ -143,6 +143,10 @@ static bool fixI8UseChain(Instruction &I,
     LoadInst *NewLoad = Builder.CreateLoad(ElementType, NewGEP);
     ReplacedValues[Load] = NewLoad;
     Load->replaceAllUsesWith(NewLoad);
+    Load->setOperand(Load->getPointerOperandIndex(),
+                     PoisonValue::get(CE->getType()));
+    if (CE->use_empty())
+      CE->destroyConstant();
     ToRemove.push_back(Load);
     return true;
   }

llvm/lib/Target/DirectX/DXILWriter/DXILBitcodeWriter.cpp

@@ -2113,7 +2113,7 @@ void DXILBitcodeWriter::writeConstants(unsigned FirstVal, unsigned LastVal,
         }
         break;
       case Instruction::GetElementPtr: {
-        Code = bitc::CST_CODE_CE_GEP;
+        Code = bitc::CST_CODE_CE_GEP_OLD;
         const auto *GO = cast<GEPOperator>(C);
         if (GO->isInBounds())
           Code = bitc::CST_CODE_CE_INBOUNDS_GEP;

llvm/lib/Target/DirectX/DirectXIRPasses/PointerTypeAnalysis.cpp

@@ -15,32 +15,46 @@
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/Operator.h"
 
 using namespace llvm;
 using namespace llvm::dxil;
 
 namespace {
 
+Type *classifyFunctionType(const Function &F, PointerTypeMap &Map);
+
 // Classifies the type of the value passed in by walking the value's users to
 // find a typed instruction to materialize a type from.
 Type *classifyPointerType(const Value *V, PointerTypeMap &Map) {
   assert(V->getType()->isPointerTy() &&
          "classifyPointerType called with non-pointer");
+
+  // A CallInst will trigger this case, and we want to classify its Function
+  // operand as a Function rather than a generic Value.
+  if (const Function *F = dyn_cast<Function>(V))
+    return classifyFunctionType(*F, Map);
+
   auto It = Map.find(V);
   if (It != Map.end())
     return It->second;
 
   Type *PointeeTy = nullptr;
-  if (auto *Inst = dyn_cast<GetElementPtrInst>(V)) {
-    if (!Inst->getResultElementType()->isPointerTy())
-      PointeeTy = Inst->getResultElementType();
+  if (auto *GEP = dyn_cast<GEPOperator>(V)) {
+    if (!GEP->getResultElementType()->isPointerTy())
+      PointeeTy = GEP->getResultElementType();
   } else if (auto *Inst = dyn_cast<AllocaInst>(V)) {
     PointeeTy = Inst->getAllocatedType();
   } else if (auto *GV = dyn_cast<GlobalVariable>(V)) {
     PointeeTy = GV->getValueType();
   }
 
   for (const auto *User : V->users()) {
+    // This assert is here because there have been ConstantExpr GEPs with no
+    // users causing the pointer type analysis to misclassify the PointeeTy.
+    [[maybe_unused]] const Constant *C = dyn_cast<Constant>(User);
+    assert((!C || !C->use_empty()) && "A Constant should not have no uses");
+
     Type *NewPointeeTy = nullptr;
     if (const auto *Inst = dyn_cast<LoadInst>(User)) {
       NewPointeeTy = Inst->getType();
@@ -49,8 +63,8 @@ Type *classifyPointerType(const Value *V, PointerTypeMap &Map) {
       // When store value is ptr type, cannot get more type info.
       if (NewPointeeTy->isPointerTy())
         continue;
-    } else if (const auto *Inst = dyn_cast<GetElementPtrInst>(User)) {
-      NewPointeeTy = Inst->getSourceElementType();
+    } else if (const auto *GEP = dyn_cast<GEPOperator>(User)) {
+      NewPointeeTy = GEP->getSourceElementType();
     }
     if (NewPointeeTy) {
       // HLSL doesn't support pointers, so it is unlikely to get more than one
@@ -204,6 +218,9 @@ PointerTypeMap PointerTypeAnalysis::run(const Module &M) {
       for (const auto &I : B) {
         if (I.getType()->isPointerTy())
           classifyPointerType(&I, Map);
+        for (const auto &O : I.operands())
+          if (O.get()->getType()->isPointerTy())
+            classifyPointerType(O.get(), Map);
       }
     }
   }

llvm/test/tools/dxil-dis/constantexpr-gep.ll

@@ -0,0 +1,26 @@
+; RUN: llc --filetype=obj %s -o - | dxil-dis -o - | FileCheck %s
+target triple = "dxil-unknown-shadermodel6.7-library"
+
+; CHECK: [[GLOBAL:@.*]] = unnamed_addr addrspace(3) global [10 x i32] zeroinitializer, align 4
+@g = local_unnamed_addr addrspace(3) global [10 x i32] zeroinitializer, align 4
+
+define i32 @fn() #0 {
+; CHECK-LABEL:  define i32 @fn()
+; CHECK-NEXT:   [[LOAD:%.*]] = load i32, i32 addrspace(3)* getelementptr inbounds ([10 x i32], [10 x i32] addrspace(3)* [[GLOBAL]], i32 0, i32 1), align 4
+; CHECK-NEXT:   ret i32 [[LOAD]]
+;
+  %gep = getelementptr [10 x i32], ptr addrspace(3) @g, i32 0, i32 1
+  %ld = load i32, ptr addrspace(3) %gep, align 4
+  ret i32 %ld
+}
+
+define i32 @fn2() #0 {
+; CHECK-LABEL:  define i32 @fn2()
+; CHECK-NEXT:   [[LOAD:%.*]] = load i32, i32 addrspace(3)* getelementptr inbounds ([10 x i32], [10 x i32] addrspace(3)* [[GLOBAL]], i32 0, i32 2), align 4
+; CHECK-NEXT:   ret i32 [[LOAD]]
+;
+  %ld = load i32, ptr addrspace(3) getelementptr ([10 x i32], ptr addrspace(3) @g, i32 0, i32 2), align 4
+  ret i32 %ld
+}
+
+attributes #0 = { "hlsl.export" }

llvm/unittests/Target/DirectX/PointerTypeAnalysisTests.cpp

@@ -8,6 +8,7 @@
 
 #include "DirectXIRPasses/PointerTypeAnalysis.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
@@ -123,6 +124,33 @@ TEST(PointerTypeAnalysis, DiscoverGEP) {
   EXPECT_THAT(Map, Contains(Pair(IsA<GetElementPtrInst>(), I64Ptr)));
 }
 
+TEST(PointerTypeAnalysis, DiscoverConstantExprGEP) {
+  StringRef Assembly = R"(
+    @g = internal global [10 x i32] zeroinitializer
+    define i32 @test() {
+      %i = load i32, ptr getelementptr ([10 x i32], ptr @g, i64 0, i64 1)
+      ret i32 %i
+    }
+  )";
+
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto M = parseAssemblyString(Assembly, Error, Context);
+  ASSERT_TRUE(M) << "Bad assembly?";
+
+  PointerTypeMap Map = PointerTypeAnalysis::run(*M);
+  ASSERT_EQ(Map.size(), 3u);
+  Type *I32Ty = Type::getInt32Ty(Context);
+  Type *I32Ptr = TypedPointerType::get(I32Ty, 0);
+  Type *I32ArrPtr = TypedPointerType::get(ArrayType::get(I32Ty, 10), 0);
+  Type *FnTy = FunctionType::get(I32Ty, {}, false);
+
+  EXPECT_THAT(Map, Contains(Pair(IsA<GlobalVariable>(), I32ArrPtr)));
+  EXPECT_THAT(Map,
+              Contains(Pair(IsA<Function>(), TypedPointerType::get(FnTy, 0))));
+  EXPECT_THAT(Map, Contains(Pair(IsA<ConstantExpr>(), I32Ptr)));
+}
+
 TEST(PointerTypeAnalysis, TraceIndirect) {
   StringRef Assembly = R"(
     define i64 @test(ptr %p) {