Place the BlockAddress type in the address space of the containing function

[arichardson]

Summary

While this should not matter for most architectures (where the program
address space is 0), it is important for CHERI. We use address space 200
for all of our code pointers and without this change we assert in
SelectionDAG handling of BlockAddress nodes.

It is also useful for AVR: previously programs targeting
AVR that attempt to read their own machine code
via a pointer to a label would instead read from RAM
using a pointer relative to the the start of program flash.

[bjope]

I don't know much about the BlockAddress concept. The LangRef says things like "always has an i8* type" and "this may be passed around as an opaque pointer sized value". But I guess it would be weird if the size doesn't match the size of pointers in the program address space, so the patch makes sense to me.

I assume that this can't be reproduced for any in-tree target?
If you can't find an in-tree reproducer, then maybe you can describe the problem a little bit more instead. Such as which assert you hit, and maybe a small stack trace. That might help when trying to motivate this patch in the future.

[arichardson]

In D48803#1149383, @bjope wrote:

I don't know much about the BlockAddress concept. The LangRef says things like "always has an i8* type" and "this may be passed around as an opaque pointer sized value". But I guess it would be weird if the size doesn't match the size of pointers in the program address space, so the patch makes sense to me.

I assume that this can't be reproduced for any in-tree target?
If you can't find an in-tree reproducer, then maybe you can describe the problem a little bit more instead. Such as which assert you hit, and maybe a small stack trace. That might help when trying to motivate this patch in the future.

Yes I'm not sure I can make a test for this with any of the existing targets. I'll see if I can get something with AVR since that sets program address space to 1.

[dylanmckay]

Nice patch, this looks useful!

Yes I'm not sure I can make a test for this with any of the existing targets. I'll see if I can get something with AVR since that sets program address space to 1.

Here's a test for you that does it.

There's another bug in LLParser that stops nonzero program address spaces from working; if the function referenced in a block address is not known in the first pass of the LLParser (for example, when the blockaddress exists earlier in the IR file than the function definition, the LLParser must insert a forward reference for the function. It does this by creating a new global variable, but it unconditionally left the global variable in the default address space of zero.

The diff I have included has a fix for this.

I've also amended the LangRef docs so that they would be accurate under the new patch.

diff --git a/docs/LangRef.rst b/docs/LangRef.rst
index 06e092fb9fc..deac223d1a1 100644
--- a/docs/LangRef.rst
+++ b/docs/LangRef.rst
@@ -3275,7 +3275,16 @@ Addresses of Basic Blocks
 ``blockaddress(@function, %block)``
 
 The '``blockaddress``' constant computes the address of the specified
-basic block in the specified function, and always has an ``i8*`` type.
+basic block in the specified function.
+
+It always has an ``i8 addrspace(P)*`` type, where ``P`` is the program
+memory address space specified in the data layout. For targets that place
+code and data in the same address space (Von-Neumann architectures) a block
+address will have the same address space as data pointers, usually
+``addrspace(0)``. Block addresses on targets that have different data and
+code address spaces (Harvard architectures) will always be in the program
+memory address space specified in the target's data layout.
+
 Taking the address of the entry block is illegal.
 
 This value only has defined behavior when used as an operand to the
diff --git a/lib/AsmParser/LLParser.cpp b/lib/AsmParser/LLParser.cpp
index 5fe1e125d48..6581436c20f 100644
--- a/lib/AsmParser/LLParser.cpp
+++ b/lib/AsmParser/LLParser.cpp
@@ -3154,9 +3154,13 @@ bool LLParser::ParseValID(ValID &ID, PerFunctionState *PFS) {
                                               std::map<ValID, GlobalValue *>()))
               .first->second.insert(std::make_pair(std::move(Label), nullptr))
               .first->second;
-      if (!FwdRef)
+      if (!FwdRef) {
         FwdRef = new GlobalVariable(*M, Type::getInt8Ty(Context), false,
-                                    GlobalValue::InternalLinkage, nullptr, "");
+                                  GlobalValue::InternalLinkage, nullptr, "",
+                                  nullptr, GlobalValue::NotThreadLocal,
+                                  M->getDataLayout().getProgramAddressSpace());
+      }
+
       ID.ConstantVal = FwdRef;
       ID.Kind = ValID::t_Constant;
       return false;
diff --git a/test/CodeGen/AVR/block-address-is-in-progmem-space.ll b/test/CodeGen/AVR/block-address-is-in-progmem-space.ll
new file mode 100644
index 00000000000..8e6e3a71062
--- /dev/null
+++ b/test/CodeGen/AVR/block-address-is-in-progmem-space.ll
@@ -0,0 +1,51 @@
+; RUN: llc -mcpu=atmega328 < %s -march=avr | FileCheck %s
+
+; This test verifies that the pointer to a basic block
+; should always be a pointer in address space 1.
+;
+; If this were not the case, then programs targeting
+; AVR that attempted to read their own machine code
+; via a pointer to a label would actually read from RAM
+; using a pointer relative to the the start of program flash.
+;
+; This would cause a load of uninitialized memory, not even
+; touching the program's machine code as otherwise desired.
+
+target datalayout = "e-P1-p:16:8-i8:8-i16:8-i32:8-i64:8-f32:8-f64:8-n8-a:8"
+
+; CHECK-LABEL: load_with_no_forward_reference
+define i8 @load_with_no_forward_reference(i8 %a, i8 %b) {
+second:
+  ; CHECK:      ldi r30, .Ltmp0+2
+  ; CHECK-NEXT: ldi r31, .Ltmp0+4
+  ; CHECK: lpm r24, Z
+  %bar = load i8, i8 addrspace(1)* blockaddress(@function_with_no_forward_reference, %second)
+  ret i8 %bar
+}
+
+; CHECK-LABEL: load_from_local_label
+define i8 @load_from_local_label(i8 %a, i8 %b) {
+entry:
+  %result1 = add i8 %a, %b
+
+  br label %second
+
+; CHECK-LABEL: .Ltmp1:
+second:
+  ; CHECK:      ldi r30, .Ltmp1+2
+  ; CHECK-NEXT: ldi r31, .Ltmp1+4
+  ; CHECK-NEXT: lpm r24, Z
+  %result2 = load i8, i8 addrspace(1)* blockaddress(@load_from_local_label, %second)
+  ret i8 %result2
+}
+
+; A function with no forward reference, right at the end
+; of the file.
+define i8 @function_with_no_forward_reference(i8 %a, i8 %b) {
+entry:
+  %result = add i8 %a, %b
+  br label %second
+second:
+  ret i8 0
+}
+

[arichardson]

Rebase on latest master and merge suggestions

[arichardson]

clang-format

[theraven]

Looks good to me. This bakes in the assumption that function pointers and basic block addresses are always in the same address space. That seems reasonable to me but it might be worth documenting in the DataLayout docs about the program address space.

[arichardson]

@dylanmckay does this change look good to you now?

[arsenm]

llvm/lib/AsmParser/LLParser.cpp

3538

Why wouldn't this come from the parent function? You should be able to mix functions with different address spaces in the same module

[bjope]

llvm/lib/AsmParser/LLParser.cpp

3538

(Maybe @arichardson got a different reason, but sharing my point-of-view here anyway.)

While it's possible to annotate calls and functions definitions with non-zero program address spaces, I think one need to be consistent. I don't think we really support multiple program address spaces (is there an actual use case for supporting that?).

I'm also not exactly sure what you mean by "parent function". The addrspce in the resulting pointer type need to match the addrspace of the function referenced in the first argument of the blockaddress. And that function has not been defined yet, since we are inside the "!F" clause.

I guess we have to trust the datalayout if the function hasn't been defined yet (or use some kind of forward ref and backtrack to fill in addrspace to get the correct type later). I wonder if we'd get some kind of type error later if we assume that datalayout is correct here, and we find a different addrspace when finding the function definition later?

(We also got the usual problem that if datalayout is set by a datalayout definition that comes later in the ll file we haven't parsed the datalayout yet. But if I remember correclty that is a general problem also for the function definitions etc.)

[dylanmckay]

llvm/lib/AsmParser/LLParser.cpp

3538

I'm also not exactly sure what you mean by "parent function". The addrspce in the resulting pointer type need to match the addrspace of the function referenced in the first argument of the blockaddress. And that function has not been defined yet, since we are inside the "!F" clause.

I suspect @arsenm is suggesting that the address space be copied from the LLParserPerFunctionState* PFS argument this function has.

For example, replace M->getDataLayout().getProgramAddressSpace() with PFS->getFunction()->getFunctionType()->getPointerAddressSpace()

• [PerFunctionState::getFunction documentation](https://ldhldh.myds.me:10081/docs/llvm700/classllvm_1_1_l_l_parser_1_1_per_function_state.html#a9b14016c1937d715c8f305742547764e)

This seems like a better alternative, as it means that if a function did opt to use a different address space from the program address space in the data layout, the blockaddresses within it will use that same address space rather than assuming the one from the datalayout.

This also makes an assumption about the address space of the target function, although I feel the case for that assumption is stronger than the current one of "assume address space from datalayout" as the address space from the parent function could be considered "closer to the source". Indeed, we directly lookup the actual address space for this block address in this branch as this is the case for when the actual function has not been defined yet and so the address space information is not available.

Please make this replacement and then the patch should be good to go.

[arichardson]

llvm/lib/AsmParser/LLParser.cpp

3538

Thanks, will do that and add a test case.

[bjope]

llvm/lib/AsmParser/LLParser.cpp

3538

Well, taking the address space from the wrong function does not help for "You should be able to mix functions with different address spaces in the same module.".

If heading in that direction then please add some nifty code comment explaining what is going on. When using the datalayout from the module it is easy to understand that we aren't using any function specific information, but if you use the function pointer from the wrong function to derive the address space it might fool someone that it either is the correct function that is being used, or that it is an "unintended bug" rather than a "hacky workaround".

[arsenm]

llvm/lib/AsmParser/LLParser.cpp

3538

block address is always in the context of a function, there's no wrong function to choose. You just use the address space from the parent function

[bjope]

llvm/lib/AsmParser/LLParser.cpp

3538

Sure, so then we agree that we don't mix address spaces for program (there can be only one and taking it from any function would be ok).

Still, there are at least some lit test cases (for example the AVR/brind.ll test modified in this patch) that have global variables involving blockaddress with forward references to functions not yet defined. In such cases there is no parent function, right? We only got the function referenced in the blockaddress argument, and its definition might not have been parsed yet. Or maybe I've simply misunderstood when we end up in this part of the code (if all the test cases pass I guess things are in order).

[arsenm]

llvm/lib/AsmParser/LLParser.cpp

3538

If the function type already exists, its address space is set. It's difficult to go back and rewrite IR types, so I don't think the IR parser would be doing this

[arichardson]

• Make use of per-function state and add various tests

[arichardson]

• Restore Constants.cpp change that was accidentally dropped while updating this revision

[arichardson]

• Update langref

[dstenb]

Hi, @arichardson! What is the status for this patch?

[arichardson]

I was waiting for an approval from one of the reviewers since it has changed quite a bit since the initial version.

[arichardson]

rebased. @dylanmckay does this look okay now?

[arsenm]

llvm/docs/LangRef.rst

4089–4090

I think this is overexplaining it. The IR is the same regardless of what the target wants to do. It should match the code address space of the function

llvm/lib/AsmParser/LLParser.cpp

3540

This should not depend on pointee element types

3542–3543

Pointee types do not really exist, the error should not mention them

[arichardson]

llvm/docs/LangRef.rst

4089–4090

I believe the following should be sufficient, right?

It always has an `i8 addrspace(P)* type, where P` is the address space
of the function containing `%block`.

@dylanmckay are you happy with dropping the the remaining text?

[arichardson]

• update error message
• Simplify langref

llvm/lib/AsmParser/LLParser.cpp

3540

I've kept this check for now (with a TODO) and only updated the error message.

[arichardson]

ping?

[arichardson]

ping?

[arichardson]

@dylanmckay ping? Would like to get this committed soon, it's only been 1.5 years

[arichardson]

@arsenm / @dylanmckay is this okay to commit now?

[arichardson]

I will commit this based on the previous approval in 1 month (1st June) unless there are any further comments.

[lebedev.ri]

In D48803#2721948, @arichardson wrote:

I will commit this based on the previous approval in 1 month (1st June) unless there are any further comments.

Looks like this wasn't re-approved after the changes were requested.
I would suggest submitting an RFC to llvm-dev.

[arichardson]

In D48803#2722120, @lebedev.ri wrote:

In D48803#2721948, @arichardson wrote:

I will commit this based on the previous approval in 1 month (1st June) unless there are any further comments.

Looks like this wasn't re-approved after the changes were requested.
I would suggest submitting an RFC to llvm-dev.

That makes sense since the pinging here has been ignored.

[bjope]

FWIW, we have merged this downstream some time ago. If I recall correctly it was related to adding support for inline asm goto.

So the solution seem to work fine for us (but have thought that it would be rude to approve it with @dylanmckay and @arsenm being ping:ed in person).

[arsenm]

llvm/lib/AsmParser/LLParser.cpp

3540

No, just rip out the type check. It's not needed now, and if it is, it's the verifier's responsibility to check it

3551

Typos programm adress

3552–3554

I prefer to not add code that you can't test and should be unreachable

[arichardson]

Thanks for the review, will update shortly.

llvm/lib/AsmParser/LLParser.cpp

3540

Will move check to verifier if there isn't one already.

3552–3554

Sounds good, will change to llvm_unreachable("").

[arichardson]

• Drop check for i8 as suggested by @arsenm

[dylanmckay]

My apologies for the couple month latency, almost every time I've checked this patch in the last few years there isn't any new activity. It looks like there's one comment from Matt last month that still needs to be addressed (remove a piece of error handling). Once updated, I'd like to get this approved.

llvm/lib/AsmParser/LLParser.cpp

3540

The error

+          if (!ExpectedTy->isPointerTy())
+            return error(ID.Loc,
+                         "type of blockaddress must be a pointer and not '" +
+                             getTypeString(ExpectedTy) + "'");

still exists and hasn't been moved to the verifier or removed from this file.

[arichardson]

I believe all remaining issues were addressed in the latest version.

llvm/lib/AsmParser/LLParser.cpp

3540

The problem was checking the pointee type , which I've removed. We need to check if it's a pointer since otherwise we can't call ExpectedTy->getPointerAddressSpace().

[dylanmckay]

Thanks for the clarification, makes sense. LGTM

llvm/docs/LangRef.rst

 .. _blockaddress:
 
 Addresses of Basic Blocks
 -------------------------
 
 ``blockaddress(@function, %block)``
 
 The '``blockaddress``' constant computes the address of the specified
-basic block in the specified function, and always has an ``i8*`` type.
+basic block in the specified function.
+
+It always has an ``i8 addrspace(P)*`` type, where ``P`` is the address space
+of the function containing ``%block`` (usually ``addrspace(0)``).
+
 Taking the address of the entry block is illegal.
 
 This value only has defined behavior when used as an operand to the
 ':ref:`indirectbr <i_indirectbr>`' or ':ref:`callbr <i_callbr>`'instruction, or
 for comparisons against null. Pointer equality tests between labels addresses
 results in undefined behavior --- though, again, comparison against null is ok,
 and no label is equal to the null pointer. This may be passed around as an
 opaque pointer sized value as long as the bits are not inspected. This
 allows ``ptrtoint`` and arithmetic to be performed on these values so
 long as the original value is reconstituted before the ``indirectbr`` or
 ``callbr`` instruction.
 
 Finally, some targets may provide defined semantics when using the value

llvm/include/llvm/AsmParser/LLParser.h

private:
     bool
     parseOptionalOperandBundles(SmallVectorImpl<OperandBundleDef> &BundleList,
                                 PerFunctionState &PFS);
 
     bool parseExceptionArgs(SmallVectorImpl<Value *> &Args,
                             PerFunctionState &PFS);
 
     // Constant Parsing.
-    bool parseValID(ValID &ID, PerFunctionState *PFS = nullptr);
+    bool parseValID(ValID &ID, PerFunctionState *PFS,
+                    Type *ExpectedTy = nullptr);
     bool parseGlobalValue(Type *Ty, Constant *&C);
     bool parseGlobalTypeAndValue(Constant *&V);
     bool parseGlobalValueVector(SmallVectorImpl<Constant *> &Elts,
                                 Optional<unsigned> *InRangeOp = nullptr);
     bool parseOptionalComdat(StringRef GlobalName, Comdat *&C);
     bool parseMetadataAsValue(Value *&V, PerFunctionState &PFS);
     bool parseValueAsMetadata(Metadata *&MD, const Twine &TypeMsg,
                               PerFunctionState *PFS);

llvm/lib/AsmParser/LLParser.cpp

if (Lex.getKind() != lltok::kw_bitcast &&
       Lex.getKind() != lltok::kw_getelementptr &&
       Lex.getKind() != lltok::kw_addrspacecast &&
       Lex.getKind() != lltok::kw_inttoptr) {
     if (parseGlobalTypeAndValue(Aliasee))
       return true;
   } else {
     // The bitcast dest type is not present, it is implied by the dest type.
     ValID ID;
-    if (parseValID(ID))
+    if (parseValID(ID, /*PFS=*/nullptr))
       return true;
     if (ID.Kind != ValID::t_Constant)
       return error(AliaseeLoc, "invalid aliasee");
     Aliasee = ID.ConstantVal;
   }
 
   Type *AliaseeType = Aliasee->getType();
   auto *PTy = dyn_cast<PointerType>(AliaseeType);

 //===----------------------------------------------------------------------===//
 
 /// parseValID - parse an abstract value that doesn't necessarily have a
 /// type implied.  For example, if we parse "4" we don't know what integer type
 /// it has.  The value will later be combined with its type and checked for
 /// sanity.  PFS is used to convert function-local operands of metadata (since
 /// metadata operands are not just parsed here but also converted to values).
 /// PFS can be null when we are not parsing metadata values inside a function.
-bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS) {
+bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
   ID.Loc = Lex.getLoc();
   switch (Lex.getKind()) {
   default:
     return tokError("expected value token");
   case lltok::GlobalID:  // @42
     ID.UIntVal = Lex.getUIntVal();
     ID.Kind = ValID::t_GlobalID;
     break;
bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
 
   case lltok::kw_blockaddress: {
     // ValID ::= 'blockaddress' '(' @foo ',' %bar ')'
     Lex.Lex();
 
     ValID Fn, Label;
 
     if (parseToken(lltok::lparen, "expected '(' in block address expression") ||
-        parseValID(Fn) ||
+        parseValID(Fn, PFS) ||
         parseToken(lltok::comma,
                    "expected comma in block address expression") ||
-        parseValID(Label) ||
+        parseValID(Label, PFS) ||
         parseToken(lltok::rparen, "expected ')' in block address expression"))
       return true;
 
     if (Fn.Kind != ValID::t_GlobalID && Fn.Kind != ValID::t_GlobalName)
       return error(Fn.Loc, "expected function name in blockaddress");
     if (Label.Kind != ValID::t_LocalID && Label.Kind != ValID::t_LocalName)
       return error(Label.Loc, "expected basic block name in blockaddress");
 
case lltok::kw_blockaddress: {
     if (!F) {
       // Make a global variable as a placeholder for this reference.
       GlobalValue *&FwdRef =
           ForwardRefBlockAddresses.insert(std::make_pair(
                                               std::move(Fn),
                                               std::map<ValID, GlobalValue *>()))
               .first->second.insert(std::make_pair(std::move(Label), nullptr))
               .first->second;
-      if (!FwdRef)
-        FwdRef = new GlobalVariable(*M, Type::getInt8Ty(Context), false,
-                                    GlobalValue::InternalLinkage, nullptr, "");
+      if (!FwdRef) {
+        unsigned FwdDeclAS;
+        if (ExpectedTy) {
+          // If we know the type that the blockaddress is being assigned to,
+          // we can use the address space of that type.
+          if (!ExpectedTy->isPointerTy())
+            return error(ID.Loc,
+                         "type of blockaddress must be a pointer and not '" +
+                             getTypeString(ExpectedTy) + "'");
+          FwdDeclAS = ExpectedTy->getPointerAddressSpace();
+        } else if (PFS) {
+          // Otherwise, we default the address space of the current function.
+          FwdDeclAS = PFS->getFunction().getAddressSpace();
+        } else {
+          llvm_unreachable("Unknown address space for blockaddress");
+        }
+        FwdRef = new GlobalVariable(
+            *M, Type::getInt8Ty(Context), false, GlobalValue::InternalLinkage,
+            nullptr, "", nullptr, GlobalValue::NotThreadLocal, FwdDeclAS);
+      }
+
       ID.ConstantVal = FwdRef;
       ID.Kind = ValID::t_Constant;
       return false;
     }
 
     // We found the function; now find the basic block.  Don't use PFS, since we
     // might be inside a constant expression.
     BasicBlock *BB;
bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
   }
 
   case lltok::kw_dso_local_equivalent: {
     // ValID ::= 'dso_local_equivalent' @foo
     Lex.Lex();
 
     ValID Fn;
 
-    if (parseValID(Fn))
+    if (parseValID(Fn, PFS))
       return true;
 
     if (Fn.Kind != ValID::t_GlobalID && Fn.Kind != ValID::t_GlobalName)
       return error(Fn.Loc,
                    "expected global value name in dso_local_equivalent");
 
     // Try to find the function (but skip it if it's forward-referenced).
     GlobalValue *GV = nullptr;
bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
   return false;
 }
 
 /// parseGlobalValue - parse a global value with the specified type.
 bool LLParser::parseGlobalValue(Type *Ty, Constant *&C) {
   C = nullptr;
   ValID ID;
   Value *V = nullptr;
-  bool Parsed = parseValID(ID) ||
+  bool Parsed = parseValID(ID, /*PFS=*/nullptr, Ty) ||
                 convertValIDToValue(Ty, ID, V, nullptr, /*IsCall=*/false);
   if (V && !(C = dyn_cast<Constant>(V)))
     return error(ID.Loc, "global values must be constants");
   return Parsed;
 }
 
 bool LLParser::parseGlobalTypeAndValue(Constant *&V) {
   Type *Ty = nullptr;
bool LLParser::convertValIDToValue(Type *Ty, ValID &ID, Value *&V,
   case ValID::t_Poison:
     // FIXME: LabelTy should not be a first-class type.
     if (!Ty->isFirstClassType() || Ty->isLabelTy())
       return error(ID.Loc, "invalid type for poison constant");
     V = PoisonValue::get(Ty);
     return false;
   case ValID::t_Constant:
     if (ID.ConstantVal->getType() != Ty)
-      return error(ID.Loc, "constant expression type mismatch");
+      return error(ID.Loc, "constant expression type mismatch: got type '" +
+                               getTypeString(ID.ConstantVal->getType()) +
+                               "' but expected '" + getTypeString(Ty) + "'");
     V = ID.ConstantVal;
     return false;
   case ValID::t_ConstantStruct:
   case ValID::t_PackedConstantStruct:
     if (StructType *ST = dyn_cast<StructType>(Ty)) {
       if (ST->getNumElements() != ID.UIntVal)
         return error(ID.Loc,
                      "initializer with struct type has wrong # elements");
bool LLParser::parseConstantValue(Type *Ty, Constant *&C) {
   default:
     return error(Loc, "expected a constant value");
   }
 }
 
 bool LLParser::parseValue(Type *Ty, Value *&V, PerFunctionState *PFS) {
   V = nullptr;
   ValID ID;
-  return parseValID(ID, PFS) ||
+  return parseValID(ID, PFS, Ty) ||
          convertValIDToValue(Ty, ID, V, PFS, /*IsCall=*/false);
 }
 
 bool LLParser::parseTypeAndValue(Value *&V, PerFunctionState *PFS) {
   Type *Ty = nullptr;
   return parseType(Ty) || parseValue(Ty, V, PFS);
 }
 
for (const auto &I : Blocks->second) {
     BasicBlock *BB;
     if (BBID.Kind == ValID::t_LocalName)
       BB = getBB(BBID.StrVal, BBID.Loc);
     else
       BB = getBB(BBID.UIntVal, BBID.Loc);
     if (!BB)
       return P.error(BBID.Loc, "referenced value is not a basic block");
 
-    GV->replaceAllUsesWith(BlockAddress::get(&F, BB));
+    Value *ResolvedVal = BlockAddress::get(&F, BB);
+    ResolvedVal = P.checkValidVariableType(BBID.Loc, BBID.StrVal, GV->getType(),
+                                           ResolvedVal, false);
+    if (!ResolvedVal)
+      return true;
+    GV->replaceAllUsesWith(ResolvedVal);
     GV->eraseFromParent();
   }
 
   P.ForwardRefBlockAddresses.erase(Blocks);
   return false;
 }
 
 /// parseFunctionBody
bool LLParser::parseInvoke(Instruction *&Inst, PerFunctionState &PFS) {
   ValID CalleeID;
   SmallVector<ParamInfo, 16> ArgList;
   SmallVector<OperandBundleDef, 2> BundleList;
 
   BasicBlock *NormalBB, *UnwindBB;
   if (parseOptionalCallingConv(CC) || parseOptionalReturnAttrs(RetAttrs) ||
       parseOptionalProgramAddrSpace(InvokeAddrSpace) ||
       parseType(RetType, RetTypeLoc, true /*void allowed*/) ||
-      parseValID(CalleeID) || parseParameterList(ArgList, PFS) ||
+      parseValID(CalleeID, &PFS) || parseParameterList(ArgList, PFS) ||
       parseFnAttributeValuePairs(FnAttrs, FwdRefAttrGrps, false,
                                  NoBuiltinLoc) ||
       parseOptionalOperandBundles(BundleList, PFS) ||
       parseToken(lltok::kw_to, "expected 'to' in invoke") ||
       parseTypeAndBasicBlock(NormalBB, PFS) ||
       parseToken(lltok::kw_unwind, "expected 'unwind' in invoke") ||
       parseTypeAndBasicBlock(UnwindBB, PFS))
     return true;
bool LLParser::parseCallBr(Instruction *&Inst, PerFunctionState &PFS) {
   LocTy RetTypeLoc;
   ValID CalleeID;
   SmallVector<ParamInfo, 16> ArgList;
   SmallVector<OperandBundleDef, 2> BundleList;
 
   BasicBlock *DefaultDest;
   if (parseOptionalCallingConv(CC) || parseOptionalReturnAttrs(RetAttrs) ||
       parseType(RetType, RetTypeLoc, true /*void allowed*/) ||
-      parseValID(CalleeID) || parseParameterList(ArgList, PFS) ||
+      parseValID(CalleeID, &PFS) || parseParameterList(ArgList, PFS) ||
       parseFnAttributeValuePairs(FnAttrs, FwdRefAttrGrps, false,
                                  NoBuiltinLoc) ||
       parseOptionalOperandBundles(BundleList, PFS) ||
       parseToken(lltok::kw_to, "expected 'to' in callbr") ||
       parseTypeAndBasicBlock(DefaultDest, PFS) ||
       parseToken(lltok::lsquare, "expected '[' in callbr"))
     return true;
 
if (TCK != CallInst::TCK_None &&
                  "expected 'tail call', 'musttail call', or 'notail call'"))
     return true;
 
   FastMathFlags FMF = EatFastMathFlagsIfPresent();
 
   if (parseOptionalCallingConv(CC) || parseOptionalReturnAttrs(RetAttrs) ||
       parseOptionalProgramAddrSpace(CallAddrSpace) ||
       parseType(RetType, RetTypeLoc, true /*void allowed*/) ||
-      parseValID(CalleeID) ||
+      parseValID(CalleeID, &PFS) ||
       parseParameterList(ArgList, PFS, TCK == CallInst::TCK_MustTail,
                          PFS.getFunction().isVarArg()) ||
       parseFnAttributeValuePairs(FnAttrs, FwdRefAttrGrps, false, BuiltinLoc) ||
       parseOptionalOperandBundles(BundleList, PFS))
     return true;
 
   // If RetType is a non-function pointer type, then this is the short syntax
   // for the call, which means that RetType is just the return type.  Infer the

 ///   ::= 'uselistorder_bb' @foo ',' %bar ',' UseListOrderIndexes
 bool LLParser::parseUseListOrderBB() {
   assert(Lex.getKind() == lltok::kw_uselistorder_bb);
   SMLoc Loc = Lex.getLoc();
   Lex.Lex();
 
   ValID Fn, Label;
   SmallVector<unsigned, 16> Indexes;
-  if (parseValID(Fn) ||
+  if (parseValID(Fn, /*PFS=*/nullptr) ||
       parseToken(lltok::comma, "expected comma in uselistorder_bb directive") ||
-      parseValID(Label) ||
+      parseValID(Label, /*PFS=*/nullptr) ||
       parseToken(lltok::comma, "expected comma in uselistorder_bb directive") ||
       parseUseListOrderIndexes(Indexes))
     return true;
 
   // Check the function.
   GlobalValue *GV;
   if (Fn.Kind == ValID::t_GlobalName)
     GV = M->getNamedValue(Fn.StrVal);

llvm/lib/IR/Constants.cpp

BlockAddress *BlockAddress::get(Function *F, BasicBlock *BB) {
   if (!BA)
     BA = new BlockAddress(F, BB);
 
   assert(BA->getFunction() == F && "Basic block moved between functions");
   return BA;
 }
 
 BlockAddress::BlockAddress(Function *F, BasicBlock *BB)
-: Constant(Type::getInt8PtrTy(F->getContext()), Value::BlockAddressVal,
-           &Op<0>(), 2) {
+    : Constant(Type::getInt8PtrTy(F->getContext(), F->getAddressSpace()),
+               Value::BlockAddressVal, &Op<0>(), 2) {
   setOperand(0, F);
   setOperand(1, BB);
   BB->AdjustBlockAddressRefCount(1);
 }
 
 BlockAddress *BlockAddress::lookup(const BasicBlock *BB) {
   if (!BB->hasAddressTaken())
     return nullptr;

llvm/test/Bitcode/blockaddress-addrspace.ll

+; RUN: rm -rf %t && split-file %s %t
+; RUN: llvm-as %t/global-use-good.ll -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/global-use-bad.ll -o /dev/null 2>&1 | FileCheck %t/global-use-bad.ll
+; RUN: llvm-as %t/global-fwddecl-good.ll -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/global-fwddecl-bad.ll -o /dev/null 2>&1 | FileCheck %t/global-fwddecl-bad.ll
+; RUN: llvm-as %t/return-fwddecl-good.ll  -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/return-fwddecl-bad.ll -o /dev/null 2>&1 | FileCheck %t/return-fwddecl-bad.ll
+; RUN: llvm-as %t/return-self-good.ll  -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/return-self-bad.ll -o /dev/null 2>&1 | FileCheck %t/return-self-bad.ll
+; RUN: not llvm-as %t/return-self-bad-2.ll -o /dev/null 2>&1 | FileCheck %t/return-self-bad-2.ll
+; RUN: not llvm-as %t/return-unknown-fn-bad.ll -o /dev/null 2>&1 | FileCheck %t/return-unknown-fn-bad.ll
+; RUN: llvm-as %t/call-fwddecl-good.ll  -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/call-fwddecl-bad.ll -o /dev/null 2>&1 | FileCheck %t/call-fwddecl-bad.ll
+; RUN: llvm-as %t/phi-good.ll  -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/phi-bad.ll -o /dev/null 2>&1 | FileCheck %t/phi-bad.ll
+; RUN: llvm-as %t/fwddecl-phi-good.ll  -o - | llvm-dis -o /dev/null
+; RUN: not llvm-as %t/fwddecl-phi-bad.ll -o /dev/null 2>&1 | FileCheck %t/fwddecl-phi-bad.ll
+; RUN: not llvm-as %t/bad-type-not-ptr.ll -o /dev/null 2>&1 | FileCheck %t/bad-type-not-ptr.ll
+; RUN: not llvm-as %t/bad-type-not-i8-ptr.ll -o /dev/null 2>&1 | FileCheck %t/bad-type-not-i8-ptr.ll
+
+
+;--- global-use-good.ll
+target datalayout = "P2"
+define void @fn_in_prog_as_implicit() {
+  unreachable
+bb:
+  ret void
+}
+define void @fn_in_prog_as_explicit() addrspace(2) {
+  unreachable
+bb:
+  ret void
+}
+define void @fn_in_other_as() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+@global1 = constant i8 addrspace(2)* blockaddress(@fn_in_prog_as_implicit, %bb)
+@global2 = constant i8 addrspace(2)* blockaddress(@fn_in_prog_as_explicit, %bb)
+@global3 = constant i8 addrspace(1)* blockaddress(@fn_in_other_as, %bb)
+
+;--- global-use-bad.ll
+define void @fn() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+@global1 = constant i8 addrspace(2)* blockaddress(@fn, %bb)
+; CHECK: [[#@LINE-1]]:38: error: constant expression type mismatch: got type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+
+; Check that a global blockaddress of a forward-declared function
+; uses the type of the global variable address space for the forward declaration
+;--- global-fwddecl-good.ll
+@global = constant i8 addrspace(2)* blockaddress(@fwddecl_in_prog_as, %bb)
+define void @fwddecl_in_prog_as() addrspace(2) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- global-fwddecl-bad.ll
+; This forward declaration does not match the actual function type so we should get an error:
+@global = constant i8 addrspace(2)* blockaddress(@fwddecl_in_unexpected_as, %bb)
+; CHECK: [[#@LINE-1]]:77: error: 'bb' defined with type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+define void @fwddecl_in_unexpected_as() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+
+; When returning blockaddresses of forward-declared functions we
+; can also use the type of the variable.
+;--- return-fwddecl-good.ll
+define i8 addrspace(2)* @take_as2() {
+  ret i8 addrspace(2)* blockaddress(@fwddecl_as2, %bb)
+}
+define i8 addrspace(1)* @take_as1() {
+  ret i8 addrspace(1)* blockaddress(@fwddecl_as1, %bb)
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+define void @fwddecl_as2() addrspace(2) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- return-fwddecl-bad.ll
+define i8 addrspace(2)* @take_bad() {
+  ret i8 addrspace(2)* blockaddress(@fwddecl_as1, %bb)
+  ; CHECK: [[#@LINE-1]]:51: error: 'bb' defined with type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- return-self-good.ll
+target datalayout = "P2"
+define i8 addrspace(0)* @take_self_as0() addrspace(0) {
+L1:
+  br label %L2
+L2:
+  ret i8 addrspace(0)* blockaddress(@take_self_as0, %L3)
+L3:
+  unreachable
+}
+define i8 addrspace(2)* @take_self_prog_as() {
+L1:
+  br label %L2
+L2:
+  ret i8 addrspace(2)* blockaddress(@take_self_prog_as, %L3)
+L3:
+  unreachable
+}
+define i8 addrspace(1)* @take_self_as1() addrspace(1) {
+L1:
+  br label %L2
+L2:
+  ret i8 addrspace(1)* blockaddress(@take_self_as1, %L3)
+L3:
+  unreachable
+}
+define i8 addrspace(2)* @take_self_as2() addrspace(2) {
+L1:
+  br label %L2
+L2:
+  ret i8 addrspace(2)* blockaddress(@take_self_as2, %L3)
+L3:
+  unreachable
+}
+
+;--- return-self-bad.ll
+target datalayout = "P2"
+define i8 addrspace(2)* @take_self_bad() addrspace(1) {
+L1:
+  br label %L2
+L2:
+  ret i8 addrspace(2)* blockaddress(@take_self_bad, %L3)
+  ; CHECK: [[#@LINE-1]]:24: error: constant expression type mismatch: got type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+L3:
+  unreachable
+}
+;--- return-self-bad-2.ll
+target datalayout = "P2"
+define i8* @take_self_bad_prog_as() {
+L1:
+  br label %L2
+L2:
+  ret i8* blockaddress(@take_self_bad_prog_as, %L3)
+  ; CHECK: [[#@LINE-1]]:11: error: constant expression type mismatch: got type 'i8 addrspace(2)*' but expected 'i8*'
+L3:
+  unreachable
+}
+
+;--- return-unknown-fn-bad.ll
+target datalayout = "P2"
+define i8 addrspace(1)* @return_unknown_fn() addrspace(1) {
+  ret i8 addrspace(1)* blockaddress(@undefined, %bb)
+  ; CHECK: [[#@LINE-1]]:37: error: expected function name in blockaddress
+}
+
+
+;--- call-fwddecl-good.ll
+target datalayout = "P2"
+define void @call_from_fn_in_as2() addrspace(2) {
+  call addrspace(2) void bitcast (i8 addrspace(2)* blockaddress(@fwddecl_as2, %bb) to void () addrspace(2)*)()
+  ret void
+}
+define void @call_from_fn_in_as1() addrspace(1) {
+  call addrspace(1) void bitcast (i8 addrspace(1)* blockaddress(@fwddecl_as1, %bb) to void () addrspace(1)*)()
+  ret void
+}
+define void @fwddecl_as2() addrspace(2) {
+  unreachable
+bb:
+  ret void
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- call-fwddecl-bad.ll
+target datalayout = "P2"
+define void @call_from_fn_in_as2_explicit() addrspace(2) {
+  call addrspace(2) void bitcast (i8 addrspace(2)* blockaddress(@fwddecl_as1, %bb) to void () addrspace(2)*)()
+  ; CHECK: [[#@LINE-1]]:79: error: 'bb' defined with type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+  ret void
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- phi-good.ll
+target datalayout = "P2"
+define i8 addrspace(1)* @f1() addrspace(1) {
+L1:
+  br label %L3
+L2:
+  br label %L3
+L3:
+  %p = phi i8 addrspace(1)* [ blockaddress(@f1, %L4), %L2 ], [ null, %L1 ]
+  ret i8 addrspace(1)* %p
+L4:
+  unreachable
+}
+define i8 addrspace(2)* @f2() {
+L1:
+  br label %L3
+L2:
+  br label %L3
+L3:
+  %p = phi i8 addrspace(2)* [ blockaddress(@f2, %L4), %L2 ], [ null, %L1 ]
+  ret i8 addrspace(2)* %p
+L4:
+  unreachable
+}
+
+;--- phi-bad.ll
+target datalayout = "P2"
+define i8* @f() {
+L1:
+  br label %L3
+L2:
+  br label %L3
+L3:
+  %p = phi i8* [ blockaddress(@f, %L4), %L2 ], [ null, %L1 ]
+  ; CHECK: [[#@LINE-1]]:18: error: constant expression type mismatch: got type 'i8 addrspace(2)*' but expected 'i8*'
+  ret i8* %p
+}
+
+; A blockaddress function forward-declaration used in a phi node should
+; create the forward declaration in the same address space as the current function
+;--- fwddecl-phi-good.ll
+define i8 addrspace(1)* @f() addrspace(1) {
+L1:
+  br label %L3
+L2:
+  br label %L3
+L3:
+  %p = phi i8 addrspace(1)* [ blockaddress(@fwddecl_as1, %bb), %L2 ], [ null, %L1 ]
+  ret i8 addrspace(1)* %p
+L4:
+  unreachable
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- fwddecl-phi-bad.ll
+define i8 addrspace(2)* @f() addrspace(2) {
+L1:
+  br label %L3
+L2:
+  br label %L3
+L3:
+  %p = phi i8 addrspace(2)* [ blockaddress(@fwddecl_as1, %bb), %L2 ], [ null, %L1 ]
+  ; CHECK: [[#@LINE-1]]:58: error: 'bb' defined with type 'i8 addrspace(1)*' but expected 'i8 addrspace(2)*'
+  ret i8 addrspace(2)* %p
+L4:
+  unreachable
+}
+define void @fwddecl_as1() addrspace(1) {
+  unreachable
+bb:
+  ret void
+}
+
+;--- bad-type-not-ptr.ll
+@global = constant i8 blockaddress(@unknown_fn, %bb)
+; CHECK: [[#@LINE-1]]:23: error: type of blockaddress must be a pointer and not 'i8'
+;--- bad-type-not-i8-ptr.ll
+@global = constant i32* blockaddress(@unknown_fn, %bb)
+; CHECK: [[#@LINE-1]]:25: error: constant expression type mismatch: got type 'i8*' but expected 'i32*'

llvm/test/CodeGen/AVR/block-address-is-in-progmem-space.ll

+; RUN: llc -mcpu=atmega328 < %s -march=avr | FileCheck %s
+
+; This test verifies that the pointer to a basic block
+; should always be a pointer in address space 1.
+;
+; If this were not the case, then programs targeting
+; AVR that attempted to read their own machine code
+; via a pointer to a label would actually read from RAM
+; using a pointer relative to the the start of program flash.
+;
+; This would cause a load of uninitialized memory, not even
+; touching the program's machine code as otherwise desired.
+
+target datalayout = "e-P1-p:16:8-i8:8-i16:8-i32:8-i64:8-f32:8-f64:8-n8-a:8"
+
+; CHECK-LABEL: load_with_no_forward_reference
+define i8 @load_with_no_forward_reference(i8 %a, i8 %b) {
+second:
+  ; CHECK:      ldi r30, .Ltmp0+2
+  ; CHECK-NEXT: ldi r31, .Ltmp0+4
+  ; CHECK: lpm r24, Z
+  %bar = load i8, i8 addrspace(1)* blockaddress(@function_with_no_forward_reference, %second)
+  ret i8 %bar
+}
+
+; CHECK-LABEL: load_from_local_label
+define i8 @load_from_local_label(i8 %a, i8 %b) {
+entry:
+  %result1 = add i8 %a, %b
+
+  br label %second
+
+; CHECK-LABEL: .Ltmp1:
+second:
+  ; CHECK:      ldi r30, .Ltmp1+2
+  ; CHECK-NEXT: ldi r31, .Ltmp1+4
+  ; CHECK-NEXT: lpm r24, Z
+  %result2 = load i8, i8 addrspace(1)* blockaddress(@load_from_local_label, %second)
+  ret i8 %result2
+}
+
+; A function with no forward reference, right at the end
+; of the file.
+define i8 @function_with_no_forward_reference(i8 %a, i8 %b) {
+entry:
+  %result = add i8 %a, %b
+  br label %second
+second:
+  ret i8 0
+}
+

llvm/test/CodeGen/AVR/brind.ll

 ; RUN: llc -mattr=sram,eijmpcall < %s -march=avr -verify-machineinstrs | FileCheck %s
 
-@brind.k = private unnamed_addr constant [2 x i8*] [i8* blockaddress(@brind, %return), i8* blockaddress(@brind, %b)], align 1
+@brind.k = private unnamed_addr constant [2 x i8 addrspace(1)*] [i8 addrspace(1)* blockaddress(@brind, %return), i8 addrspace(1)* blockaddress(@brind, %b)], align 1
 
 define i8 @brind(i8 %p) {
 ; CHECK-LABEL: brind:
 ; CHECK: ijmp
 entry:
   %idxprom = sext i8 %p to i16
-  %arrayidx = getelementptr inbounds [2 x i8*], [2 x i8*]* @brind.k, i16 0, i16 %idxprom
-  %s = load i8*, i8** %arrayidx
-  indirectbr i8* %s, [label %return, label %b]
+  %arrayidx = getelementptr inbounds [2 x i8 addrspace(1)*], [2 x i8 addrspace(1)*]* @brind.k, i16 0, i16 %idxprom
+  %s = load i8 addrspace(1)*, i8 addrspace(1)** %arrayidx
+  indirectbr i8 addrspace(1)* %s, [label %return, label %b]
 b:
   br label %return
 return:
   %retval.0 = phi i8 [ 4, %b ], [ 2, %entry ]
   ret i8 %retval.0
 }