[AArch64][PAC] Select MOVK for ptrauth.blend intrinsic.

[ab]

Summary

In a sense, this is implementing an ABI decision (how to lower the software construct of "blend"), but if there are interesting variants to consider, this could be made object-file-format-specific in some way.

[kristof.beyls]

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll

19

Overall, I think this patch looks fine.
I notice that the definition of the llvm.ptrauth.blend intrinisic is (per https://llvm.org/docs/PointerAuth.html#llvm-ptrauth-blend):

The integer discriminator argument is a small integer, as specified by the target.

and the signature of the llvm.ptrauth.blend is documented as:

declare i64 @llvm.ptrauth.blend(i64 <address discriminator>, i64 <integer discriminator>)

It makes me wonder what should happen when an immediate too large to fit into an i16 is passed....
I'm guessing with the patch as currently written, it would result in the backend failing to do instruction selection?
Or rather DAGISel failing to do instruction selection and GlobalISel succeeding, taking the lower 16 bits of the immediate?

I wonder if for consistency reasons it would be easily possible to make DAGISel behave like GlobalISel and accept immediates larger than 16 bits, and take the lower 16 bit value?
That would be behavior that presumably is consistent also with when the blend value comes from a register rather than an immediate?

[ab]

• further constrain GISel immediate blends to fail if wider than 16 bits

[ab]

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll

19

Interesting! I would go the other way, and consider the GISel behavior too lax. I added a check that makes it fail isel as well. It's true we don't specify the intrinsics strictly enough, but for AArch64 I'd consider it always incorrect to use a wider value.

It's also true that that's the emergent behavior for non-immediate blends, but I'd argue those are by far the exception: I think the only legitimate use is a dynamic loader that really does get the discriminator values from memory. Anything else is a red flag (and I considered somehow hiding them in the compiler, or detecting accidental uses, but that's tricky.)

How does that sound?

[kristof.beyls]

I had just 2 more minor nitpicks - mostly about adding 2 more regression tests to test/document expected behavior in corner cases.

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8409

It seems that the blend intrinsic can be implemented on cores that do not implement the PAuth extension. Which made me wonder if it should be allowed or not for a user to use the blend intrinsic when targeting a core not supporting PAuth?
I don't have a very strong argument either way. I wonder if you happen to have an opinion?
One thing I did think off - but couldn't answer from just looking at this patch - are DAGISel and GISel consistent in rejecting the blend intrinsic when PAuth is not implemented?
I think it would make sense to have a regression test checking that the blend intrinsic is indeed accepted/rejected as expected when targeting a core that does not support PAuth.

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll

19

That sound good to me!
Shouldn't there also be a regression test that verifies that instruction selection fails for an immediate that does not fit in 16 bits?

[ab]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8409

It seems that the blend intrinsic can be implemented on cores that do not implement the PAuth extension. Which made me wonder if it should be allowed or not for a user to use the blend intrinsic when targeting a core not supporting PAuth?

I don't have a strong argument either, but until we have a use for it on other cores, I'd leave it as disallowed there?
Somewhat related: XPAC/strip is, however, commonly needed on cores that don't support PAC (or aren't even aarch64), usually for debugging and crash analysis tools. But that one actually depends on aarch64 processor state (TBI, T0Sz), sadly.

One thing I did think off - but couldn't answer from just looking at this patch - are DAGISel and GISel consistent in rejecting the blend intrinsic when PAuth is not implemented?

They're not, and I found out by later tests breaking ;) On SDAG, if we can't match the 16-bit immediate, we'll fallback to the arbitrary GPR64 value (which isn't ideal by itself.) So we can't easily fail on wider values, at least not without some machinery to match this case (maybe matching any immediate and erroring out when trying to encode; I'm not sure we can fail in a good way at isel time)

[kristof.beyls]

llvm/lib/Target/AArch64/AArch64InstrInfo.td

8409

It seems that the blend intrinsic can be implemented on cores that do not implement the PAuth extension. Which made me wonder if it should be allowed or not for a user to use the blend intrinsic when targeting a core not supporting PAuth?

I don't have a strong argument either, but until we have a use for it on other cores, I'd leave it as disallowed there?

That seems OK to me.

One thing I did think off - but couldn't answer from just looking at this patch - are DAGISel and GISel consistent in rejecting the blend intrinsic when PAuth is not implemented?

They're not, and I found out by later tests breaking ;) On SDAG, if we can't match the 16-bit immediate, we'll fallback to the arbitrary GPR64 value (which isn't ideal by itself.) So we can't easily fail on wider values, at least not without some machinery to match this case (maybe matching any immediate and erroring out when trying to encode; I'm not sure we can fail in a good way at isel time)

I also don't see an immediate solution to this.
It seems to me that we shouldn't hold up progress on the upstreaming of this functionality on this minor issue. In other words, I think this issue can be ignored, at least for now.

[atrosinenko]

Could this patch be landed? (I can commit it on your behalf if appropriate)

Overall, I would like to assist in preparing for review and upstreaming more PAC-related patches from https://github.com/ahmedbougacha/llvm-project/commits/eng/arm64e-upstream-llvmorg .

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

5875–5876

[nitpick] Maybe reorder these lines so the first removeOperand doesn't affect the index passed to the second one

[kristof.beyls]

I think this patch can be landed. @ab, do you see any reason not to land this patch?

[atrosinenko]

I am planning to land this patch the next week, if there will be no objections.

[ab]

In D132384#4405529, @atrosinenko wrote:

I am planning to land this patch the next week, if there will be no objections.

I don’t think that’s how we do things. Either way I was out for a while and will update the patches; thanks for taking a look.

llvm/lib/Target/AArch64/AArch64InstrFormats.td

 let ParserMatchClass = AsmImmRange<0, 65535>, PrintMethod = "printImmHex" in {
 def timm32_0_65535 : Operand<i32>, TImmLeaf<i32, [{
   return ((uint32_t)Imm) < 65536;
 }]>;
 
 def timm64_0_65535 : Operand<i64>, TImmLeaf<i64, [{
   return ((uint64_t)Imm) < 65536;
 }]>;
+
+def imm64_0_65535 : Operand<i64>, ImmLeaf<i64, [{
+  return ((uint64_t)Imm) < 65536;
+}]>;
 }
 
 // imm0_255 predicate - True if the immediate is in the range [0,255].
 def Imm0_255Operand : AsmImmRange<0,255>;
 
 def imm0_255 : Operand<i32>, ImmLeaf<i32, [{
   return ((uint32_t)Imm) < 256;
 }]> {

llvm/lib/Target/AArch64/AArch64InstrInfo.td

let Predicates = [HasMOPS], Defs = [NZCV], Size = 12, mayStore = 1 in {
   }
 }
 let Predicates = [HasMOPS, HasMTE], Defs = [NZCV], Size = 12, mayLoad = 0, mayStore = 1 in {
   def MOPSMemorySetTaggingPseudo : Pseudo<(outs GPR64common:$Rd_wb, GPR64:$Rn_wb),
                                           (ins GPR64common:$Rd, GPR64:$Rn, GPR64:$Rm),
                                           [], "$Rd = $Rd_wb,$Rn = $Rn_wb">, Sched<[]>;
 }
 
+//-----------------------------------------------------------------------------
+// v8.3 Pointer Authentication late patterns
+
+let Predicates = [HasPAuth] in {
+def : Pat<(int_ptrauth_blend GPR64:$Rd, imm64_0_65535:$imm),
+          (MOVKXi GPR64:$Rd, (trunc_imm imm64_0_65535:$imm), 48)>;
+def : Pat<(int_ptrauth_blend GPR64:$Rd, GPR64:$Rn),
+          (BFMXri GPR64:$Rd, GPR64:$Rn, 16, 15)>;
+}
+
+//-----------------------------------------------------------------------------
+
 // This gets lowered into an instruction sequence of 20 bytes
 let Defs = [X16, X17], mayStore = 1, isCodeGenOnly = 1, Size = 20 in
 def StoreSwiftAsyncContext
       : Pseudo<(outs), (ins GPR64:$ctx, GPR64sp:$base, simm9:$offset),
                []>, Sched<[]>;
 
 def AArch64AssertZExtBool : SDNode<"AArch64ISD::ASSERT_ZEXT_BOOL", SDT_assert>;
 def : Pat<(AArch64AssertZExtBool GPR32:$op),
           (i32 GPR32:$op)>;
 
 include "AArch64InstrAtomics.td"
 include "AArch64SVEInstrInfo.td"
 include "AArch64SMEInstrInfo.td"
 include "AArch64InstrGISel.td"

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

case Intrinsic::ptrauth_strip: {
 
     MIB.buildInstr(Opcode, {DstReg}, {ValReg});
 
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);
     RBI.constrainGenericRegister(ValReg, AArch64::GPR64RegClass, MRI);
     I.eraseFromParent();
     return true;
   }
+  case Intrinsic::ptrauth_blend: {
+    MachineFunction &MF = *I.getParent()->getParent();
+    auto RHS = getIConstantVRegVal(I.getOperand(3).getReg(), MRI);
+    if (RHS) {
+      I.setDesc(TII.get(AArch64::MOVKXi));
+      I.removeOperand(1);
+      I.removeOperand(2);
+      MachineInstrBuilder(MF, I)
+          .addImm(RHS->getZExtValue() & 0xffff)
+          .addImm(48)
+          .constrainAllUses(TII, TRI, RBI);
+    } else {
+      I.setDesc(TII.get(AArch64::BFMXri));
+      I.removeOperand(1);
+      MachineInstrBuilder(MF, I).addImm(16).addImm(15).constrainAllUses(
+          TII, TRI, RBI);
+    }
+    return true;
+  }
   case Intrinsic::frameaddress:
   case Intrinsic::returnaddress: {
     MachineFunction &MF = *I.getParent()->getParent();
     MachineFrameInfo &MFI = MF.getFrameInfo();
 
     unsigned Depth = I.getOperand(2).getImm();
     Register DstReg = I.getOperand(0).getReg();
     RBI.constrainGenericRegister(DstReg, AArch64::GPR64RegClass, MRI);

llvm/test/CodeGen/AArch64/ptrauth-intrinsic-blend.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=0 | FileCheck %s
+; RUN: llc < %s -mtriple arm64e-apple-darwin -verify-machineinstrs -global-isel=1 -global-isel-abort=1 | FileCheck %s
+
+define i64 @test_blend(i64 %arg, i64 %arg1) {
+; CHECK-LABEL: test_blend:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    bfi x0, x1, #48, #16
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 %arg, i64 %arg1)
+  ret i64 %tmp
+}
+
+define i64 @test_blend_constant(i64 %arg) {
+; CHECK-LABEL: test_blend_constant:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    movk x0, #12345, lsl #48
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 %arg, i64 12345)
+  ret i64 %tmp
+}
+
+; Blend isn't commutative.
+define i64 @test_blend_constant_swapped(i64 %arg) {
+; CHECK-LABEL: test_blend_constant_swapped:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov w8, #12345
+; CHECK-NEXT:    bfi x8, x0, #48, #16
+; CHECK-NEXT:    mov x0, x8
+; CHECK-NEXT:    ret
+  %tmp = call i64 @llvm.ptrauth.blend(i64 12345, i64 %arg)
+  ret i64 %tmp
+}
+
+declare i64 @llvm.ptrauth.blend(i64, i64)