[llvm] [SPIRV] Add legalization pass for zero-size arrays (PR #172367)

[Nick Sarnie via llvm-commits]

================
@@ -0,0 +1,359 @@
+//===- SPIRVLegalizeZeroSizeArrays.cpp - Legalize zero-size arrays -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// SPIR-V does not support zero-size arrays unless it is within a shader. This
+// pass legalizes zero-size arrays ([0 x T]) in unsupported cases.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SPIRVLegalizeZeroSizeArrays.h"
+#include "SPIRV.h"
+#include "SPIRVTargetMachine.h"
+#include "SPIRVUtils.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/InstVisitor.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "spirv-legalize-zero-size-arrays"
+
+using namespace llvm;
+
+namespace {
+
+bool hasZeroSizeArray(const Type *Ty) {
+  if (const ArrayType *ArrTy = dyn_cast<ArrayType>(Ty)) {
+    if (ArrTy->getNumElements() == 0)
+      return true;
+    return hasZeroSizeArray(ArrTy->getElementType());
+  }
+
+  if (const StructType *StructTy = dyn_cast<StructType>(Ty)) {
+    for (Type *ElemTy : StructTy->elements()) {
+      if (hasZeroSizeArray(ElemTy))
+        return true;
+    }
+  }
+
+  return false;
+}
+
+bool shouldLegalizeInstType(const Type *Ty) {
+  if (const ArrayType *ArrTy = dyn_cast_if_present<ArrayType>(Ty)) {
+    return ArrTy->getNumElements() == 0 ||
+           shouldLegalizeInstType(ArrTy->getElementType());
+  }
+  return false;
+}
+
+class SPIRVLegalizeZeroSizeArraysImpl
+    : public InstVisitor<SPIRVLegalizeZeroSizeArraysImpl> {
+  friend class InstVisitor<SPIRVLegalizeZeroSizeArraysImpl>;
+
+public:
+  SPIRVLegalizeZeroSizeArraysImpl(const SPIRVTargetMachine *TM)
+      : InstVisitor(), TM(TM) {}
+  bool runOnModule(Module &M);
+
+  // TODO: Handle GEP, PHI.
+  void visitAllocaInst(AllocaInst &AI);
+  void visitLoadInst(LoadInst &LI);
+  void visitStoreInst(StoreInst &SI);
+  void visitSelectInst(SelectInst &Sel);
+  void visitExtractValueInst(ExtractValueInst &EVI);
+  void visitInsertValueInst(InsertValueInst &IVI);
+
+private:
+  Type *legalizeType(Type *Ty);
+  Constant *legalizeConstant(Constant *C);
+
+  const SPIRVTargetMachine *TM;
+  DenseMap<Type *, Type *> TypeMap;
+  DenseMap<GlobalVariable *, GlobalVariable *> GlobalMap;
+  SmallVector<Instruction *, 16> ToErase;
+  bool Modified = false;
+};
+
+class SPIRVLegalizeZeroSizeArraysLegacy : public ModulePass {
+public:
+  static char ID;
+  SPIRVLegalizeZeroSizeArraysLegacy(const SPIRVTargetMachine *TM)
+      : ModulePass(ID), TM(TM) {}
+  StringRef getPassName() const override {
+    return "SPIRV Legalize Zero-Size Arrays";
+  }
+  bool runOnModule(Module &M) override {
+    SPIRVLegalizeZeroSizeArraysImpl Impl(TM);
+    return Impl.runOnModule(M);
+  }
+
+private:
+  const SPIRVTargetMachine *TM;
+};
+
+Type *SPIRVLegalizeZeroSizeArraysImpl::legalizeType(Type *Ty) {
+  auto It = TypeMap.find(Ty);
+  if (It != TypeMap.end())
+    return It->second;
+
+  Type *LegalizedTy = Ty;
+
+  if (ArrayType *ArrTy = dyn_cast<ArrayType>(Ty)) {
+    if (ArrTy->getNumElements() == 0) {
+      LegalizedTy = PointerType::get(
+          Ty->getContext(),
+          storageClassToAddressSpace(SPIRV::StorageClass::Generic));
+    } else if (Type *ElemTy = legalizeType(ArrTy->getElementType());
+               ElemTy != ArrTy->getElementType()) {
+      LegalizedTy = ArrayType::get(ElemTy, ArrTy->getNumElements());
+    }
+  } else if (StructType *StructTy = dyn_cast<StructType>(Ty)) {
+    SmallVector<Type *, 8> ElemTypes;
+    bool Changed = false;
+    for (Type *ElemTy : StructTy->elements()) {
+      Type *LegalizedElemTy = legalizeType(ElemTy);
+      ElemTypes.push_back(LegalizedElemTy);
+      Changed |= LegalizedElemTy != ElemTy;
+    }
+    if (Changed) {
+      LegalizedTy =
+          StructTy->hasName()
+              ? StructType::create(StructTy->getContext(), ElemTypes,
+                                   (StructTy->getName() + ".legalized").str(),
+                                   StructTy->isPacked())
+              : StructType::get(StructTy->getContext(), ElemTypes,
+                                StructTy->isPacked());
+    }
+  }
+
+  TypeMap[Ty] = LegalizedTy;
+  return LegalizedTy;
+}
+
+Constant *SPIRVLegalizeZeroSizeArraysImpl::legalizeConstant(Constant *C) {
+  if (!C || !hasZeroSizeArray(C->getType()))
+    return C;
+
+  if (GlobalVariable *GV = dyn_cast<GlobalVariable>(C)) {
+    if (GlobalVariable *NewGV = GlobalMap.lookup(GV))
+      return NewGV;
+    return C;
+  }
+
+  Type *NewTy = legalizeType(C->getType());
+  if (isa<UndefValue>(C))
+    return PoisonValue::get(NewTy);
+  if (isa<ConstantAggregateZero>(C))
+    return Constant::getNullValue(NewTy);
+
+  if (ConstantArray *CA = dyn_cast<ConstantArray>(C)) {
+    SmallVector<Constant *, 8> Elems;
+    for (Use &U : CA->operands())
+      Elems.push_back(legalizeConstant(cast<Constant>(U)));
----------------
sarnex wrote:

For infinite loops, the number of operands must be finite and isn't changed as part of the loop, so it must terminate. Thanks

https://github.com/llvm/llvm-project/pull/172367