[llvm] [SPIR-V] Add pass to fixup global variable AS (PR #124591)

[Michal Paszkowski via llvm-commits]

================
@@ -0,0 +1,593 @@
+//===-- SPIRVFixAddressSpace.cpp ----------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+// This pass is Vulkan specific as Logical SPIR-V doesn't support pointer
+// cast.
+//
+// In LLVM IR, global and local variables are by default in the same address
+// space. In SPIR-V, global and local variables live in a different address
+// space (storage class in the SPIR-V parlance).
+//
+// This means the following function cannot be lowered to SPIR-V:
+//   static int global = 0;
+//   int& return_ref(int& ref, bool select) {
+//    return select ? ref : global;
+//   }
+//
+// A solution is to force inline, but this would prevent us from emitting SPIR-V
+// libraries. Another solution is to move all globals to local variables, but
+// this also blocks libraries. The last solution is to replace all local
+// variables with global variables. This is possible because Vulkan SPIR-V
+// completely forbids static recursion.
+//
+// This pass replace all alloca instruction with a new global variable.
+// In addition, it moves all such allocations into the `Private` address space.
+//
+// After this pass, no variable or pointer should reference the default address
+// space.
+//
+// Note:
+//  LLVM IR has address spaces for functions, but SPIR-V doesn't. In addition,
+//  Vulkan disallow function pointers and indirect jump, meaning we could never
+//  have a pointer storing the function address. For this reason, functions are
+//  left in the default address space, but all pointer operands to the default
+//  AS are rewritten to point to the AS `Private`. This kind of blind rewrite
+//  simplifies the code, but can only work with those assumptions.
+//===----------------------------------------------------------------------===//
+
+#include "Analysis/SPIRVConvergenceRegionAnalysis.h"
+#include "SPIRV.h"
+#include "SPIRVSubtarget.h"
+#include "SPIRVTargetMachine.h"
+#include "SPIRVUtils.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/CodeGen/IntrinsicLowering.h"
+#include "llvm/IR/CFG.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/IntrinsicsSPIRV.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/NoFolder.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/PassRegistry.h"
+#include "llvm/Transforms/Utils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/LoopSimplify.h"
+#include "llvm/Transforms/Utils/LowerMemIntrinsics.h"
+#include <queue>
+#include <stack>
+#include <type_traits>
+#include <unordered_set>
+
+using namespace llvm;
+using namespace SPIRV;
+
+namespace llvm {
+void initializeSPIRVFixAddressSpacePass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+
+constexpr unsigned GlobalAddressSpace =
+    storageClassToAddressSpace(SPIRV::StorageClass::Private);
+
+// Returns true if the given type or any subtype contains a pointer to the
+// default address space.
+bool typeRequiresConversion(Type *T) {
+  if (T->isPointerTy())
+    return T->getPointerAddressSpace() == 0;
+
+  if (T->isArrayTy())
+    return typeRequiresConversion(T->getArrayElementType());
+  if (T->isStructTy()) {
+    for (unsigned I = 0; I < T->getStructNumElements(); ++I)
+      if (typeRequiresConversion(T->getStructElementType(I)))
+        return true;
+    return false;
+  }
+  if (FunctionType *FT = dyn_cast<FunctionType>(T)) {
+    if (typeRequiresConversion(FT->getReturnType()))
+      return true;
+    for (Type *TP : FT->params())
+      if (typeRequiresConversion(TP))
+        return true;
+    return false;
+  }
+
+  return false;
+}
+
+class SPIRVFixAddressSpace : public InstVisitor<SPIRVFixAddressSpace>,
+                             public ModulePass {
+
+  // Types are supposed to be unique. When using Type::get, there is a lookup to
+  // only create types on demand. StructType::get only allows creating literal
+  // structs, meaning we would lose the type name. This forces us to use
+  // StructType::create, which doesn't deduplicates. This means we must bring
+  // our own old-type/new-type map to prevent creating distinct types when we
+  // shouldn't.
+  std::unordered_map<Type *, Type *> ConvertedTypes;
----------------
michalpaszkowski wrote:

You could replace `std::unordered_map` with `llvm::DenseMap` especially for pointer to pointer lookups. I think using `std::unordered_map` [is generally discouraged](https://llvm.org/docs/ProgrammersManual.html#other-map-like-container-options).

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