[IR][Verifier] Relax restriction on alloca address spaces

[wingo]

Summary

In the WebAssembly target, we would like to allow alloca in two address
spaces. The alloca instruction already has an address space argument,
but the verifier asserts that the address space of an alloca is the
default alloca address space from the datalayout. This patch removes
this restriction. Targets that would like to impose additional
restrictions should do so via target-specific verification passes.

[wingo]

Cc arsenm, as he added the address space restriction for alloca when he added address space support for alloca in https://github.com/llvm/llvm-project/commit/3c1fc768ed6cf2b01463df036ae6ae3b1f0de632.

[wingo]

Pre-merge checks show that https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/invalid-alloca.ll fails; will need to follow up in a similar way.

[tlively]

The code looks fine, but explicitly requesting a review from @arsenm because I don't know the history of this restriction or how this affects other targets.

@wingo, have you investigated what would happen if we tried to use the normal stack address space in IR, but extracted reftype stack slots into a separate stack in the WebAssembly backend, perhaps specifically in WebAssemblyFrameLowering::emit{Prologue,Epilogue}?

[arsenm]

I did envision this being useful, however there are a few points of concern:

1. This does remove the safety check for the common case of creating an alloca without checking the datalayout's default alloca address space
2. How is this going to be represented in codegen? In SelectionDAG the pointer type is immediately discarded and frame indexes don't have a tracked address space

[wingo]

I posted https://reviews.llvm.org/D101140 just for context -- based on this patch, and @pmatos's work on https://reviews.llvm.org/D95425.

In D101045#2709494, @tlively wrote:

@wingo, have you investigated what would happen if we tried to use the normal stack address space in IR, but extracted reftype stack slots into a separate stack in the WebAssembly backend, perhaps specifically in WebAssemblyFrameLowering::emit{Prologue,Epilogue}?

Humm, good question. The issue as I understand it would be that letting reference-typed values slide all the way to frame lowering would be what to do with uses of the alloca at instruction-selection time. For i32 locals, isel lowersa (set i32:$dst (load frameindex:$idx)) to i32.load, but what would we do for externref values?

But if we took the same solution as in my downstream patch, maybe keeping the alloca in AS 0 might with, where the ElementType is our AS 1 externref representation. Though, I do wonder what risk we might be introducing with regards to optimizations / transformations thinking that the pointer has a size and can be written to linear memory.

[wingo]

Thanks for your time, arsenm.

In D101045#2710615, @arsenm wrote:

I did envision this being useful, however there are a few points of concern:

1. This does remove the safety check for the common case of creating an alloca without checking the datalayout's default alloca address space

Yeah I can see why you might want this. If we want to keep the footgun-protection, one way would be to allow a target to explicitly opt out of the alloca address-space check. Something like:

target datalayout = "A1*"

to indicate that this target uses AS 1 by default for alloca, but that it can use others. Then we'd adapt the Verifier check to see if this datalayout allows alloca in other address spaces. WDYT?

1. How is this going to be represented in codegen? In SelectionDAG the pointer type is immediately discarded and frame indexes don't have a tracked address space

We override TargetLowering::getPointerTy(DataLayout&, unsigned AS) to return a different MVT for these address spaces. We then have custom lowering for ISD::LOAD / ISD::STORE for those special MVTs.

Just for context, on the wasm target, each function has two stack frames: one managed by the webassembly run-time, populated by named locals, and not addressable; and one in memory managed as on conventional targets, sharing an address space with data and heap. Clang produces IR where locals start life as alloca's as usual, and as usual the optimizer lifts most of them to SSA values in SROA. SSA values end up lowering to the "named" part of the stack. Those locals still left as allocas are put on linear memory. This patch is motivated by ongoing work to add support for a new kind of value that can't be written to linear memory. We are currently representing these values as opaque pointers in non-default address spaces, and have some special lowering hooks for these. We also use some special address spaces to denote restricted pointers to these values.

[wingo]

Instead of globally relaxing the restriction on alloca address spaces,
allow targets to choose whether to impose the restriction via the
datalayout.

[arsenm]

It hadn't occurred to me to add a field to the datalayout purely for the verifier. That seems a bit strange, since it's not exactly a target property.

What is the plan to use this in codegen?

llvm/docs/LangRef.rst

9566

c/designed/designated/?

[arsenm]

In D101045#2711261, @wingo wrote:

Just for context, on the wasm target, each function has two stack frames: one managed by the webassembly run-time, populated by named locals, and not addressable; and one in memory managed as on conventional targets, sharing an address space with data and heap. Clang produces IR where locals start life as alloca's as usual, and as usual the optimizer lifts most of them to SSA values in SROA. SSA values end up lowering to the "named" part of the stack. Those locals still left as allocas are put on linear memory. This patch is motivated by ongoing work to add support for a new kind of value that can't be written to linear memory. We are currently representing these values as opaque pointers in non-default address spaces, and have some special lowering hooks for these. We also use some special address spaces to denote restricted pointers to these values.

But are these still tracked in the MachineFrameInfo like a regular frame index?

[wingo]

In D101045#2711957, @arsenm wrote:

It hadn't occurred to me to add a field to the datalayout purely for the verifier. That seems a bit strange, since it's not exactly a target property.

Yeah understood. If we want some targets to have the check and some to not have the check, do you know of a better place to express this? I thought the same place that the target's alloca address space is defined might not be unreasonable, but I am new here :)

In D101045#2711959, @arsenm wrote:

In D101045#2711261, @wingo wrote:

Just for context, on the wasm target, each function has two stack frames: one managed by the webassembly run-time, populated by named locals, and not addressable; and one in memory managed as on conventional targets, sharing an address space with data and heap. Clang produces IR where locals start life as alloca's as usual, and as usual the optimizer lifts most of them to SSA values in SROA. SSA values end up lowering to the "named" part of the stack. Those locals still left as allocas are put on linear memory. This patch is motivated by ongoing work to add support for a new kind of value that can't be written to linear memory. We are currently representing these values as opaque pointers in non-default address spaces, and have some special lowering hooks for these. We also use some special address spaces to denote restricted pointers to these values.

But are these still tracked in the MachineFrameInfo like a regular frame index?

They reach the MachineFrameInfo, but only ephemerally. At the beginning of isel, the static alloca's get collected and the MachineFrameInfo populated with abstract objects and frame indexes. Then when isel sees a load / store to these objects, the first load/store of the alloca causes it to be lifted out from the MachineFrameInfo and instead allocates it on the target-specific stack (as a WebAssembly local), and the object in the MachineFrameInfo is marked dead. Loads and stores to these objects lower to accesses to the parallel stack instead of to objects in MachineFrameInfo.

Obviously this doesn't allow for by-reference access to these allocations. This is OK for the WebAssembly target because these opaque types have source-level restrictions similar to the ones in place for ARM's SVE values -- see the ARM C Language Extensions (ACLE), §3.2.1 https://developer.arm.com/documentation/100987/0000.

[wingo]

• Fix typo in LangRef

[wingo]

Address feedback from tlively

[wingo]

Fix spurious whitespace

[arsenm]

In D101045#2718774, @wingo wrote:

In D101045#2711957, @arsenm wrote:

It hadn't occurred to me to add a field to the datalayout purely for the verifier. That seems a bit strange, since it's not exactly a target property.

Yeah understood. If we want some targets to have the check and some to not have the check, do you know of a better place to express this? I thought the same place that the target's alloca address space is defined might not be unreasonable, but I am new here :)

I think we should have target specific IR verifiers, not pollute the datalayout. There are a number of other places such a verifier would be useful

[wingo]

In D101045#2755666, @arsenm wrote:

In D101045#2718774, @wingo wrote:

In D101045#2711957, @arsenm wrote:

It hadn't occurred to me to add a field to the datalayout purely for the verifier. That seems a bit strange, since it's not exactly a target property.

Yeah understood. If we want some targets to have the check and some to not have the check, do you know of a better place to express this? I thought the same place that the target's alloca address space is defined might not be unreasonable, but I am new here :)

I think we should have target specific IR verifiers, not pollute the datalayout. There are a number of other places such a verifier would be useful

Sure but this would be a much larger design effort. As I see it I have three options right now:

1. Remove the restriction imposed on alloca with an explicitly specified address space
2. Keep the restriction but allow datalayout to specify multiple valid alloca address spaces (this patch)
3. Add a target-specific IR verifier somehow. It couldn't extend Verifier because Verifier is private; would have to be a separate pass. However surely in this case we'd remove the restriction on alloca address spaces in the core verifier.

I am leaning back towards (1), wdyt @arsenm ?

[arsenm]

In D101045#2762675, @wingo wrote:

I am leaning back towards (1), wdyt @arsenm ?

Yes, just remove the restriction. The verifier is a separate issue

[wingo]

Go back to first version of patch, removing verifier restriction on alloca address spaces

[wingo]

Fix drop-debug-info-nonzero-alloca test

[arsenm]

llvm/docs/LangRef.rst

9526

I would rephrase this instead of dropping the statement altogether. Without target knowledge, any generic code introducing an alloca should be using the alloca address space in the datalayout

[wingo]

Tweak alloca langref language

[wingo]

llvm/docs/LangRef.rst

9526

Thanks for feedback; updated. I was unsure how precisely to say "alloca outside the stack address space may be meaningless", given that this is the langref; thoughts welcome if there is a better way to say this.

[wingo]

Remove AMDGPU test that alloca in wrong addrspace fails to verify

[arsenm]

LGTM

[tlively]

Thanks for your review, @arsenm! LGTM from me as well.

llvm/docs/LangRef.rst

 ``G<address space>``
     Specifies the address space to be used by default when creating global
     variables. If omitted, the globals address space defaults to the default
     address space 0.
     Note: variable declarations without an address space are always created in
     address space 0, this property only affects the default value to be used
     when creating globals without additional contextual information (e.g. in
     LLVM passes).
-``A<address space>``
+``A<address space>:<restricted>``
     Specifies the address space of objects created by '``alloca``'.
-    Defaults to the default address space of 0.
+    Defaults to the default address space of 0. The ``:<restricted>``
+    suffix is optional; possible values are:
+
+    * ``r``: Restricted ``alloca`` address spaces. The verifier will
+      check that ``alloca`` instructions are in the specified address
+      space. This is the default.
+    * ``u``: Unrestricted ``alloca`` address spaces. Though the default
+      address space for ``alloca`` is ``<address space>``, the verifier will
+      allow ``alloca`` in other address spaces.
 ``p[n]:<size>:<abi>:<pref>:<idx>``
     This specifies the *size* of a pointer and its ``<abi>`` and
     ``<pref>``\erred alignments for address space ``n``. The fourth parameter
     ``<idx>`` is a size of index that used for address calculation. If not
     specified, the default index size is equal to the pointer size. All sizes
     are in bits. The address space, ``n``, is optional, and if not specified,
     denotes the default address space 0. The value of ``n`` must be
     in the range [1,2^23).
::
 
       <result> = alloca [inalloca] <type> [, <ty> <NumElements>] [, align <alignment>] [, addrspace(<num>)]     ; yields type addrspace(num)*:result
 
 Overview:
 """""""""
 
 The '``alloca``' instruction allocates memory on the stack frame of the
 currently executing function, to be automatically released when this
-function returns to its caller. The object is always allocated in the
-address space for allocas indicated in the datalayout.
+function returns to its caller.
 
 Arguments:
 """"""""""
 
 The '``alloca``' instruction allocates ``sizeof(<type>)*NumElements``
 bytes of memory on the runtime stack, returning a pointer of the
 appropriate type to the program. If "NumElements" is specified, it is
 the number of elements allocated, otherwise "NumElements" is defaulted

 space for the allocation.'``alloca``'d memory is automatically released
 when the function returns. The '``alloca``' instruction is commonly used
 to represent automatic variables that must have an address available. When
 the function returns (either with the ``ret`` or ``resume`` instructions),
 the memory is reclaimed. Allocating zero bytes is legal, but the returned
 pointer may not be unique. The order in which memory is allocated (ie.,
 which way the stack grows) is not specified.
 
+If not explicitly specified, the address space of the result defaults to
+the designated alloca address space for the target.  Note that most
+targets require that the object be allocated in the address space for
+allocas indicated in the datalayout.  See :ref:`datalayout
+string<langref_datalayout>` for more on alloca address spaces.
+
 If the returned pointer is used by :ref:`llvm.lifetime.start <int_lifestart>`,
 the returned object is initially dead.
 See :ref:`llvm.lifetime.start <int_lifestart>` and
 :ref:`llvm.lifetime.end <int_lifeend>` for the precise semantics of
 lifetime-manipulating intrinsics.
 
 Example:
 """"""""

llvm/include/llvm/IR/DataLayout.h

enum class FunctionPtrAlignType {
     Independent,
     /// The function pointer alignment is a multiple of the function alignment.
     MultipleOfFunctionAlign,
   };
 private:
   /// Defaults to false.
   bool BigEndian;
 
+  bool RestrictedAllocaAddrSpace;
   unsigned AllocaAddrSpace;
   MaybeAlign StackNaturalAlign;
   unsigned ProgramAddrSpace;
   unsigned DefaultGlobalsAddrSpace;
 
   MaybeAlign FunctionPtrAlign;
   FunctionPtrAlignType TheFunctionPtrAlignType;
 
public:
   DataLayout(const DataLayout &DL) { *this = DL; }
 
   ~DataLayout(); // Not virtual, do not subclass this class
 
   DataLayout &operator=(const DataLayout &DL) {
     clear();
     StringRepresentation = DL.StringRepresentation;
     BigEndian = DL.isBigEndian();
+    RestrictedAllocaAddrSpace = DL.restrictedAllocaAddrSpace();
     AllocaAddrSpace = DL.AllocaAddrSpace;
     StackNaturalAlign = DL.StackNaturalAlign;
     FunctionPtrAlign = DL.FunctionPtrAlign;
     TheFunctionPtrAlignType = DL.TheFunctionPtrAlignType;
     ProgramAddrSpace = DL.ProgramAddrSpace;
     DefaultGlobalsAddrSpace = DL.DefaultGlobalsAddrSpace;
     ManglingMode = DL.ManglingMode;
     LegalIntWidths = DL.LegalIntWidths;
bool exceedsNaturalStackAlignment(Align Alignment) const {
     return StackNaturalAlign && (Alignment > *StackNaturalAlign);
   }
 
   Align getStackAlignment() const {
     assert(StackNaturalAlign && "StackNaturalAlign must be defined");
     return *StackNaturalAlign;
   }
 
+  bool restrictedAllocaAddrSpace() const { return RestrictedAllocaAddrSpace; }
   unsigned getAllocaAddrSpace() const { return AllocaAddrSpace; }
 
   /// Returns the alignment of function pointers, which may or may not be
   /// related to the alignment of functions.
   /// \see getFunctionPtrAlignType
   MaybeAlign getFunctionPtrAlign() const { return FunctionPtrAlign; }
 
   /// Return the type of function pointer alignment.

llvm/lib/IR/DataLayout.cpp

static const LayoutAlignElem DefaultAlignments[] = {
     {AGGREGATE_ALIGN, 0, Align(1), Align(8)}   // struct
 };
 
 void DataLayout::reset(StringRef Desc) {
   clear();
 
   LayoutMap = nullptr;
   BigEndian = false;
+  RestrictedAllocaAddrSpace = true;
   AllocaAddrSpace = 0;
   StackNaturalAlign.reset();
   ProgramAddrSpace = 0;
   DefaultGlobalsAddrSpace = 0;
   FunctionPtrAlign.reset();
   TheFunctionPtrAlignType = FunctionPtrAlignType::Independent;
   ManglingMode = MM_None;
   NonIntegralAddressSpaces.clear();
while (!Desc.empty()) {
     case 'P': { // Function address space.
       if (Error Err = getAddrSpace(Tok, ProgramAddrSpace))
         return Err;
       break;
     }
     case 'A': { // Default stack/alloca address space.
       if (Error Err = getAddrSpace(Tok, AllocaAddrSpace))
         return Err;
+      if (!Rest.empty()) {
+        if (Error Err = split(Rest, ':', Split))
+          return Err;
+        if (Tok.empty())
+          return reportError("Unexpected trailing ':' in alloca addrspace "
+                             "specifier in datalayout string");
+        if (Tok.size() > 1)
+          return reportError("Unknown 'restricted' specifier in alloca "
+                             "addrspace specifier in datalayout string");
+        switch (Tok[0]) {
+        default:
+          return reportError("Unknown 'restricted' specifier in alloca "
+                             "addrspace specifier in datalayout string");
+        case 'r':
+          RestrictedAllocaAddrSpace = true;
+          break;
+        case 'u':
+          RestrictedAllocaAddrSpace = false;
+          break;
+        }
+      }
       break;
     }
     case 'G': { // Default address space for global variables.
       if (Error Err = getAddrSpace(Tok, DefaultGlobalsAddrSpace))
         return Err;
       break;
     }
     case 'm':

 DataLayout::DataLayout(const Module *M) {
   init(M);
 }
 
 void DataLayout::init(const Module *M) { *this = M->getDataLayout(); }
 
 bool DataLayout::operator==(const DataLayout &Other) const {
   bool Ret = BigEndian == Other.BigEndian &&
+             RestrictedAllocaAddrSpace == Other.RestrictedAllocaAddrSpace &&
              AllocaAddrSpace == Other.AllocaAddrSpace &&
              StackNaturalAlign == Other.StackNaturalAlign &&
              ProgramAddrSpace == Other.ProgramAddrSpace &&
              DefaultGlobalsAddrSpace == Other.DefaultGlobalsAddrSpace &&
              FunctionPtrAlign == Other.FunctionPtrAlign &&
              TheFunctionPtrAlignType == Other.TheFunctionPtrAlignType &&
              ManglingMode == Other.ManglingMode &&
              LegalIntWidths == Other.LegalIntWidths &&

llvm/lib/IR/Verifier.cpp

for (const User *U : SwiftErrorVal->users()) {
     if (auto *Call = dyn_cast<CallBase>(U))
       verifySwiftErrorCall(*const_cast<CallBase *>(Call), SwiftErrorVal);
   }
 }
 
 void Verifier::visitAllocaInst(AllocaInst &AI) {
   SmallPtrSet<Type*, 4> Visited;
   PointerType *PTy = AI.getType();
-  // TODO: Relax this restriction?
-  Assert(PTy->getAddressSpace() == DL.getAllocaAddrSpace(),
-         "Allocation instruction pointer not in the stack address space!",
-         &AI);
+  Assert(PTy->getAddressSpace() == DL.getAllocaAddrSpace() ||
+             !DL.restrictedAllocaAddrSpace(),
+         "Allocation instruction pointer not in the stack address space!", &AI);
   Assert(AI.getAllocatedType()->isSized(&Visited),
          "Cannot allocate unsized type", &AI);
   Assert(AI.getArraySize()->getType()->isIntegerTy(),
          "Alloca array size must have integer type", &AI);
   Assert(AI.getAlignment() <= Value::MaximumAlignment,
          "huge alignment values are unsupported", &AI);
 
   if (AI.isSwiftError()) {

llvm/test/Verifier/alloca-alternate-addrspace-invalid-0.ll

+; RUN: not llvm-as %s 2>&1 | FileCheck %s
+
+target datalayout = "A1"
+
+define void @use_alloca() {
+; CHECK: Allocation instruction pointer not in the stack address space!
+  %alloca_scalar_no_align = alloca i32, addrspace(2)
+  ret void
+}

llvm/test/Verifier/alloca-alternate-addrspace-invalid-1.ll

+; RUN: not llvm-as %s 2>&1 | FileCheck %s
+
+target datalayout = "A1:r"
+
+define void @use_alloca() {
+; CHECK: Allocation instruction pointer not in the stack address space!
+  %alloca_scalar_no_align = alloca i32, addrspace(2)
+  ret void
+}

llvm/test/Verifier/alloca-alternate-addrspace-invalid-2.ll

+; RUN: not llvm-as < %s 2>&1 | FileCheck %s
+
+target datalayout = ""
+
+define void @use_alloca() {
+; CHECK: Allocation instruction pointer not in the stack address space!
+  %alloca_scalar_no_align = alloca i32, addrspace(2)
+  ret void
+}

llvm/test/Verifier/alloca-alternate-addrspace-valid.ll

+; RUN: llvm-as %s -o /dev/null
+
+; Should assemble without error.
+
+target datalayout = "A1:u"
+
+define void @use_alloca0() {
+  %alloca_scalar_no_align = alloca i32, addrspace(2)
+  ret void
+}
+
+define void @use_alloca1() {
+  %alloca_scalar_no_align = alloca i32, align 4, addrspace(2)
+  ret void
+}
+
+define void @use_alloca2() {
+  %alloca_scalar_no_align = alloca i32, align 4, addrspace(2), !foo !0
+  ret void
+}
+
+!0 = !{}