[SPIR-V] Fix bad insertion for type/id MIR (#109686)

Keenuts · web-flow · commit 02a334de6690 · 2024-09-24T15:24:21.000+02:00
Those instructions were inserted either after the instruction using it,
or in the middle of the module.
The first directly causes an issue. The second causes a more subtle
issue: the first type the type is inserted, the emission is fine, but
the second times, the first instruction is reused, without checking its
position in the function. This can lead to the second usage dominating
the definition.

In SPIR-V, types are usually in the header, above all code definition,
but at this stage I don't think we can, so what I do instead is to emit
it in the first basic block.

This commit reduces the failed tests with expensive checks from 107 to
71.

Signed-off-by: Nathan Gauër &lt;brioche@google.com&gt;
diff --git a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.cpp
@@ -25,6 +25,7 @@
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Casting.h"
 #include <cassert>
+#include <functional>
 
 using namespace llvm;
 SPIRVGlobalRegistry::SPIRVGlobalRegistry(unsigned PointerSize)
@@ -83,8 +84,11 @@ inline Register createTypeVReg(MachineIRBuilder &MIRBuilder) {
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeBool(MachineIRBuilder &MIRBuilder) {
-  return MIRBuilder.buildInstr(SPIRV::OpTypeBool)
-      .addDef(createTypeVReg(MIRBuilder));
+
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeBool)
+        .addDef(createTypeVReg(MIRBuilder));
+  });
 }
 
 unsigned SPIRVGlobalRegistry::adjustOpTypeIntWidth(unsigned Width) const {
@@ -118,24 +122,53 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeInt(unsigned Width,
     MIRBuilder.buildInstr(SPIRV::OpCapability)
         .addImm(SPIRV::Capability::ArbitraryPrecisionIntegersINTEL);
   }
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeInt)
-                 .addDef(createTypeVReg(MIRBuilder))
-                 .addImm(Width)
-                 .addImm(IsSigned ? 1 : 0);
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeInt)
+        .addDef(createTypeVReg(MIRBuilder))
+        .addImm(Width)
+        .addImm(IsSigned ? 1 : 0);
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeFloat(uint32_t Width,
                                                MachineIRBuilder &MIRBuilder) {
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeFloat)
-                 .addDef(createTypeVReg(MIRBuilder))
-                 .addImm(Width);
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeFloat)
+        .addDef(createTypeVReg(MIRBuilder))
+        .addImm(Width);
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeVoid(MachineIRBuilder &MIRBuilder) {
-  return MIRBuilder.buildInstr(SPIRV::OpTypeVoid)
-      .addDef(createTypeVReg(MIRBuilder));
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeVoid)
+        .addDef(createTypeVReg(MIRBuilder));
+  });
+}
+
+SPIRVType *SPIRVGlobalRegistry::createOpType(
+    MachineIRBuilder &MIRBuilder,
+    std::function<MachineInstr *(MachineIRBuilder &)> Op) {
+  auto oldInsertPoint = MIRBuilder.getInsertPt();
+  MachineBasicBlock *OldMBB = &MIRBuilder.getMBB();
+
+  auto LastInsertedType = LastInsertedTypeMap.find(CurMF);
+  if (LastInsertedType != LastInsertedTypeMap.end()) {
+    MIRBuilder.setInsertPt(*MIRBuilder.getMF().begin(),
+                           LastInsertedType->second->getIterator());
+  } else {
+    MIRBuilder.setInsertPt(*MIRBuilder.getMF().begin(),
+                           MIRBuilder.getMF().begin()->begin());
+    auto Result = LastInsertedTypeMap.try_emplace(CurMF, nullptr);
+    assert(Result.second);
+    LastInsertedType = Result.first;
+  }
+
+  MachineInstr *Type = Op(MIRBuilder);
+  LastInsertedType->second = Type;
+
+  MIRBuilder.setInsertPt(*OldMBB, oldInsertPoint);
+  return Type;
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeVector(uint32_t NumElems,
@@ -147,11 +180,12 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeVector(uint32_t NumElems,
           EleOpc == SPIRV::OpTypeBool) &&
          "Invalid vector element type");
 
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeVector)
-                 .addDef(createTypeVReg(MIRBuilder))
-                 .addUse(getSPIRVTypeID(ElemType))
-                 .addImm(NumElems);
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeVector)
+        .addDef(createTypeVReg(MIRBuilder))
+        .addUse(getSPIRVTypeID(ElemType))
+        .addImm(NumElems);
+  });
 }
 
 std::tuple<Register, ConstantInt *, bool, unsigned>
@@ -688,22 +722,25 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeArray(uint32_t NumElems,
   SPIRVType *SpvTypeInt32 = getOrCreateSPIRVIntegerType(32, MIRBuilder);
   Register NumElementsVReg =
       buildConstantInt(NumElems, MIRBuilder, SpvTypeInt32, EmitIR);
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeArray)
-                 .addDef(createTypeVReg(MIRBuilder))
-                 .addUse(getSPIRVTypeID(ElemType))
-                 .addUse(NumElementsVReg);
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeArray)
+        .addDef(createTypeVReg(MIRBuilder))
+        .addUse(getSPIRVTypeID(ElemType))
+        .addUse(NumElementsVReg);
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeOpaque(const StructType *Ty,
                                                 MachineIRBuilder &MIRBuilder) {
   assert(Ty->hasName());
   const StringRef Name = Ty->hasName() ? Ty->getName() : "";
   Register ResVReg = createTypeVReg(MIRBuilder);
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeOpaque).addDef(ResVReg);
-  addStringImm(Name, MIB);
-  buildOpName(ResVReg, Name, MIRBuilder);
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeOpaque).addDef(ResVReg);
+    addStringImm(Name, MIB);
+    buildOpName(ResVReg, Name, MIRBuilder);
+    return MIB;
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(const StructType *Ty,
@@ -717,14 +754,16 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypeStruct(const StructType *Ty,
     FieldTypes.push_back(getSPIRVTypeID(ElemTy));
   }
   Register ResVReg = createTypeVReg(MIRBuilder);
-  auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeStruct).addDef(ResVReg);
-  for (const auto &Ty : FieldTypes)
-    MIB.addUse(Ty);
-  if (Ty->hasName())
-    buildOpName(ResVReg, Ty->getName(), MIRBuilder);
-  if (Ty->isPacked())
-    buildOpDecorate(ResVReg, MIRBuilder, SPIRV::Decoration::CPacked, {});
-  return MIB;
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    auto MIB = MIRBuilder.buildInstr(SPIRV::OpTypeStruct).addDef(ResVReg);
+    for (const auto &Ty : FieldTypes)
+      MIB.addUse(Ty);
+    if (Ty->hasName())
+      buildOpName(ResVReg, Ty->getName(), MIRBuilder);
+    if (Ty->isPacked())
+      buildOpDecorate(ResVReg, MIRBuilder, SPIRV::Decoration::CPacked, {});
+    return MIB;
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOrCreateSpecialType(
@@ -739,17 +778,22 @@ SPIRVType *SPIRVGlobalRegistry::getOpTypePointer(
     MachineIRBuilder &MIRBuilder, Register Reg) {
   if (!Reg.isValid())
     Reg = createTypeVReg(MIRBuilder);
-  return MIRBuilder.buildInstr(SPIRV::OpTypePointer)
-      .addDef(Reg)
-      .addImm(static_cast<uint32_t>(SC))
-      .addUse(getSPIRVTypeID(ElemType));
+
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypePointer)
+        .addDef(Reg)
+        .addImm(static_cast<uint32_t>(SC))
+        .addUse(getSPIRVTypeID(ElemType));
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeForwardPointer(
     SPIRV::StorageClass::StorageClass SC, MachineIRBuilder &MIRBuilder) {
-  return MIRBuilder.buildInstr(SPIRV::OpTypeForwardPointer)
-      .addUse(createTypeVReg(MIRBuilder))
-      .addImm(static_cast<uint32_t>(SC));
+  return createOpType(MIRBuilder, [&](MachineIRBuilder &MIRBuilder) {
+    return MIRBuilder.buildInstr(SPIRV::OpTypeForwardPointer)
+        .addUse(createTypeVReg(MIRBuilder))
+        .addImm(static_cast<uint32_t>(SC));
+  });
 }
 
 SPIRVType *SPIRVGlobalRegistry::getOpTypeFunction(
diff --git a/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h b/llvm/lib/Target/SPIRV/SPIRVGlobalRegistry.h
@@ -64,6 +64,10 @@ class SPIRVGlobalRegistry {
   SmallPtrSet<const Type *, 4> TypesInProcessing;
   DenseMap<const Type *, SPIRVType *> ForwardPointerTypes;
 
+  // Stores for each function the last inserted SPIR-V Type.
+  // See: SPIRVGlobalRegistry::createOpType.
+  DenseMap<const MachineFunction *, MachineInstr *> LastInsertedTypeMap;
+
   // if a function returns a pointer, this is to map it into TypedPointerType
   DenseMap<const Function *, TypedPointerType *> FunResPointerTypes;
 
@@ -97,6 +101,13 @@ class SPIRVGlobalRegistry {
                         SPIRV::AccessQualifier::AccessQualifier AccessQual,
                         bool EmitIR);
 
+  // Internal function creating the an OpType at the correct position in the
+  // function by tweaking the passed "MIRBuilder" insertion point and restoring
+  // it to the correct position. "Op" should be the function creating the
+  // specific OpType you need, and should return the newly created instruction.
+  SPIRVType *createOpType(MachineIRBuilder &MIRBuilder,
+                          std::function<MachineInstr *(MachineIRBuilder &)> Op);
+
 public:
   SPIRVGlobalRegistry(unsigned PointerSize);
 
diff --git a/llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp b/llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp
@@ -389,9 +389,7 @@ void processInstr(MachineInstr &MI, MachineIRBuilder &MIB,
       createNewIdReg(nullptr, MI.getOperand(0).getReg(), MRI, *GR).first;
   AssignTypeInst.getOperand(1).setReg(NewReg);
   MI.getOperand(0).setReg(NewReg);
-  MIB.setInsertPt(*MI.getParent(),
-                  (MI.getNextNode() ? MI.getNextNode()->getIterator()
-                                    : MI.getParent()->end()));
+  MIB.setInsertPt(*MI.getParent(), MI.getIterator());
   for (auto &Op : MI.operands()) {
     if (!Op.isReg() || Op.isDef())
       continue;
