Correctly legalise stackmap operands

[vext01]

Summary

Hi everyone,

This is my first real contribution to LLVM, so please be gentle!

For context, we are writing a JIT that will require LLVM's stackmap facility in order to reconstruct the native stack when deoptimising from specialised JITted code back into generic AOT-compiled code.

During my experimentation with stackmaps, I've identified a few problems that I hope to be able to fix. This diff makes a start at addressing the first of these issues, namely: https://github.com/llvm/llvm-project/issues/21657

In short, passing illegally-typed live variable operands to llvm.experimental.stackmap (at -O1 or above) will make LLVM crash by an assertion failure. The issue is that stackmap operands are emitted directly to target nodes in the selection DAG, which means their operands don't get legalised.

The diff below introduces a (non-target) stackmap DAG node, so that the stackmap node and its operands can take part in regular legalisations.

Although ninja check passes, I don't expect this to be ready to be merged just yet because:

• I only legalise integer operands and constant operands for now. I expect more needs to be done to legalise floats (and are there any other types that require legalisation?)
• GC statepoints and patchpoints have the same problem as stackmaps and will require the same treatment.
• There are a couple of questions, marked in the diff with XXX. I'd appreciate it if someone could comment on those points.

So although incomplete, I'm raising this now to check that I'm headed in the right direction. Any and all feedback would be much appreciated (even if you just say "yes, this looks correct/on-track", that's useful).

(The diff is formatted with git clang-format version 11, as found on Debian. Hope that's OK)

[vext01]

I wasn't sure who to add as reviewers, so on the sugesstion of people in IRC, I've added a few people who have recently touched selection the selection DAG. Hope that's OK.

[nikic]

Can you please upload the diff with context (-U99999)?

I'm not familiar with this intrinsic, and it looks like it isn't specified in LangRef :/ It doesn't look like the original author is still active. Adding @reames who might be familiar with JIT and deoptimization.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

494–495

I would move this above the included file (preferably in the same position as the code is defined in the enum)

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2256–2257

The opcode assert is implied by the cast. There's also probably no point to asserting it's i64

2265–2266

Ditto

2279–2280

You seem to be fixing up some Constants to TargetConstant here. It would be better to just do this up front in SelectionDAGBuilder like it was before

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

If you're checking debug output you need REQUIRES: asserts. However, I don't think checking the legalizer output is the most helpful thing here. Better to check the final output

9

Probably should add some cases with excessively wide types, vectors and FP

[vext01]

Updated diff with full context.

[vext01]

I'm not familiar with this intrinsic, and it looks like it isn't specified in LangRef.

There's a little info here:
https://llvm.org/docs/StackMaps.html

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2279–2280

These parameters probably should have been marked immarg in the intrinsic definition?

[vext01]

Responded to @arsenm's comments. Some outstanding questions.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2256–2257

If the assertion fails, do we not end up doing a bogus cast and invoking undefined behaviour which will probably lead to a crash much later? Wouldn't it be best to have the assertion crash early?

2279–2280

You seem to be fixing up some Constants to TargetConstant here. It would be better to just do this up front in SelectionDAGBuilder like it was before

Wasn't that the very reason that the constants were not being legalised? Because they were already target constants?

These parameters probably should have been marked immarg in the intrinsic definition?

Should that be handled as a separate change, since I'm not actually touching the intrinsic's definition in this work?

I don't know how the workflow works after phabricator? Do I get an opportunity to make things into neat, self-contained commits?

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

At what level should I check?

I wanted to check at the MIR level, but it doesn't show the types there.

9

Yep. I was only working on integer types for now, but I can add changes to ensure those work too.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2256–2257

The point of cast<> is it does the type assertion for you

2279–2280

I assume these are the fixed operands at the start of the argument list, and not the variadic section. I'm also assuming legalization is only relevant for the variadic arguments.

Changing the intrinsic would be a separate change. Phabricator lets you add parent/child revisions to track related changes

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

The types are meaningless after selection, so that makes sense. I would default to codegen to the end. MIR is less stable and what you care about is that the types were legalized to and selected to something, not the types themselves

[vext01]

Respond to more of @arsenm's comments. Thanks.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2256–2257

I had assumed that cast<>was an alias to static_cast<>, but now I see that LLVM defines it! I had no idea!

(forgive me, I'm pretty much learning C++ on the fly here)

So I agree that asserting the type of the opCode is redundant, but:

There's also probably no point to asserting it's i64

How comes? Isn't it good to check the sanity of the DAG we fed in? If we remove the assertion I think the stackmap code would eventually crash somewhere a lot later in the pipeline, but wouldn't it be good to crash early?

2279–2280

I assume these are the fixed operands at the start of the argument list, and not the variadic section. I'm also assuming legalization is only relevant for the variadic arguments.

Ah, do you mean we should emit the first two arguments directly to target nodes? Assuming so, that'd probably work and might be a little faster.

By the way, did you have any comments on the XXXs below?

Changing the intrinsic would be a separate change. Phabricator lets you add parent/child revisions to track related changes.

OK. Let's do that as a separate change.

[vext01]

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

I would default to codegen to the end

Do you mean to match asm code for the target architecture? Isn't that also going to be fragile?

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2279–2280

Yes, if they needed to be target constants to begin with they should start as targetconstants

2286

FrameIndex will also be materialized into a register, unlike TargetFrameIndex

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

Yes. No, this is a target specific test anyway.

[vext01]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2286

And is that correct?

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

Stackmap doesn't emit any code that we could match, but we could match the raw bytes of the .llvmbc section, just it's very fiddly.

If that's the only way, so be it.

[vext01]

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

2

sorry, not .llvmbc, I meant the stackmap section.

[reames]

General direction looks reasonable. I do want to warn you that to my knowledge, no one is using stackmap or patchpoint. I believe they've been effectively dead code for the last couple years. statepoint is used, and should provide a superset of the stackmap/patchpoint functionality.

Code structure wise, would it be possibly to split this patch? (Warning: I am no selectiondag expert!) Could we add the SDNode in an NFC patch which does not legalize, and then handle legalization in a separate patch? If we can, it might make some of the review (mapping argument orders, etc) more straight forward.

[arsenm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2286

Well the current code is emitting TargetFrameIndexes into the argument list, so I assume this should also just go direct to TargetFrameIndex from the start

[vext01]

Hi Philip,

In D125680#3519150, @reames wrote:

General direction looks reasonable.

Good!

I do want to warn you that to my knowledge, no one is using stackmap or patchpoint. I believe they've been effectively dead code for the last couple years. statepoint is used, and should provide a superset of the stackmap/patchpoint functionality.

This makes me a little nervous. For our system to succeed we need them to work 100%!

Code structure wise, would it be possibly to split this patch? (Warning: I am no selectiondag expert!) Could we add the SDNode in an NFC patch which does not legalize, and then handle legalization in a separate patch? If we can, it might make some of the review (mapping argument orders, etc) more straight forward.

We could split the patch, but I don't see much benefit. It will mean picking the diff apart later, just to have an intermediate "still crashing" state.

I don't know what the others think...

[arsenm]

I don't particularly care about splitting the patch. You can't really break this down very far, and the individual changes aren't really individually testable

[vext01]

Revised diff addressing many comments, most notably:

• Emit target nodes at DAG-construction time if we know they are already legal.
• Test raw bytes of stackmap section.

Still to do:

• i128 live vars causes a crash.
• floats
• vectors
• make llvm.experimental.patchpoint and llvm.experimental.gc.statepoint legalise properly too.

[dantrushin]

In D125680#3522332, @vext01 wrote:

• make llvm.experimental.patchpoint and llvm.experimental.gc.statepoint legalise properly too.

Please do not touch statepoint in this patch. If you do not really need it, leave it alone at all.

[vext01]

Sorry, bodged the last diff upload.

[vext01]

Please do not touch statepoint in this patch. If you do not really need it, leave it alone at all.

Just to clarify why I've proposed this:

llvm.experimental.stackmap, llvm.experimental.patchpoint and llvm.experimental.gc.statepoint used to all share the same DAG building code via a function called addStackMapLiveVars() in SelectionDAGBuilder.cpp.

This function incorrectly emits target nodes, thus erroneously side-stepping the legalisation stage.

My change (so-far) detaches only llvm.experimental.stackmap from this function.

Ideally, all three facilities would use a common function that emits target-independent nodes when legalisation is required, but if upstream is against that, we should at the very least add a big flashing neon message to the effect of:

// XXX: this is incorrect because it emits target nodes, meaning that operands do not get legalised.

Are we on the same page?

[dantrushin]

In D125680#3522346, @vext01 wrote:

Please do not touch statepoint in this patch. If you do not really need it, leave it alone at all.

Just to clarify why I've proposed this:

llvm.experimental.stackmap, llvm.experimental.patchpoint and llvm.experimental.gc.statepoint used to all share the same DAG building code via a function called addStackMapLiveVars() in SelectionDAGBuilder.cpp.

No, Statepoint does not use this function.

This function incorrectly emits target nodes, thus erroneously side-stepping the legalisation stage.

My change (so-far) detaches only llvm.experimental.stackmap from this function.

This is enough for a _single_ patch
You can (and better) do patchpoint in a separate patch.
Statepoint is much more complicated and is in active productionuse today. It is too easy to break, so we will need to negotiate first if you'll want to change it

Ideally, all three facilities would use a common function that emits target-independent nodes when legalisation is required, but if upstream is against that, we should at the very least add a big flashing neon message to the effect of:

// XXX: this is incorrect because it emits target nodes, meaning that operands do not get legalised.

For statepoint deopt arguments are always legal right now (as well as gc live pointers), so they do not need any special processing.

Are we on the same page?

As far as I know nobody cares about stackmap/patchpoint. Neither do we.
But as I said, STATEPOINT is very important for us in its current form.

[vext01]

Statepoint is much more complicated and is in active productionuse today. It is too easy to break, so we will need to negotiate first if you'll want to change it

We don't need statepoint, and if what you say is true, that it doesn't use addStackMapLiveVars(), then I have no need to touch statepoint and everyone is happy :)

[vext01]

Here is the latest version of this diff.

Also in the meantime I've extended the stackmap tests to give us a reasonable degree of confidence that this change is not breaking things.

• Added some i128 tests in D126069 (merged)
• Added some float tests in D126204 (merged)
• Added runtime stackmap tests in D126552 (needs review)

I had planned to include legalisation tests for things like vectors, structs and integers greater than 64-bits, but since those types already don't work properly with stackmaps (i.e. when already legal), I think that should be a separate change.

For more info on those issues, see:

• https://github.com/llvm/llvm-project/issues/26431#issuecomment-1132697106
• https://github.com/llvm/llvm-project/issues/55609
• https://github.com/llvm/llvm-project/issues/55649
• https://github.com/llvm/llvm-project/issues/55613

I hope this diff is nearing completion, but D126552 should go in first (I'm struggling to find reviewers for that one, any takers?).

[vext01]

Here's the same diff, just formatted correctly.

Since we are now no longer waiting on D126552, would it be OK for this to go into main please?

(I now have commit access, so with permission, I can do that myself)

[dantrushin]

In general looks OK to me (modulo few style nits).
But I'm not an ISEL type legalization expert. Would be nice to have LGTM from someone else as well.

llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp

2920 ↗

(On Diff #437863)

Nit: could you have it in single line as cases above (if it fits in 80 chars)?

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

1729

Nit: same comment for a single line

4669

Nit: and here too

4684

Nit: extra empty line

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

9423

Use existing Ops instead of introducing new vector?

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

490

Single line again

[vext01]

llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp

2920 ↗

(On Diff #437863)

I think git clang-format wanted it this way, but I will check.

[vext01]

I've just realised that no test covers the case where a i128 constant is used as a stackmap live variable.

Before any attempt to legalise stackmap operands, a stackmap call would be directly translated to a target node and any constant i128 operand becomes a TargetConstant (also of type i128). LLVM would then crash if the value of the i128 exceeds what can be expressed by an i64 (see this bug).

This change causes LLVM to try to legalise the constant i128, which comes with some other caveats. When legalising a constant i128, LLVM is going to want to expand the type into smaller chunks (in the case of a constant i128, two smaller i64s). This is problematic because the stackmap format has no way to express such an expansion -- it is assumed that a constant cannot be larger than 64-bits ([see LargeConstant here](https://llvm.org/docs/StackMaps.html#stack-map-format).

A similar issue will arise for non-constant i128s, which may be split across registers during legalisation. Again, the stackmap format has no way to express that: it is assumed that if something lives in registers, it cannot span more than a single register.

I don't want to change the stackmap format right now, so I'm going to have to try and find an imperfect solution in the interim. I hope that's OK.

Sorry this is taking so long.

[vext01]

Made changes requested by @dantrushin, fixed (and tested) the i128 issue, and tidied up quite a bit.

I was able to revert changes I had made to some of the existing assertions by having the stackmap legalisation routines replace the node themselves.

Hoping this is ready now. I have a branch to do the same for llvm.experimental.patchpoint(), which I'd like to raise a diff for, but this diff needs to go in first.

llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp

2920 ↗

(On Diff #437863)

Yep. If I fold the lines, then git clang-format just undoes it. Should I override it in this instance, or?

[arsenm]

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

84

For a general test, a non-register argument may be more helpful

85

Is it worth testing vectors?

[vext01]

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

85

We spoke on IRC, but for the others, vector support for stackmaps is currently broken, so I've made no attempt to address them for now.

https://github.com/llvm/llvm-project/issues/55613

[vext01]

To ease reviewing, here's what changed since the last version:

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index 64d92e60f5de..7ed2808600cf 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -5515,8 +5515,12 @@ SDValue DAGTypeLegalizer::ExpandIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
         DAG.getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
     NewOps.push_back(DAG.getTargetConstant(CN->getZExtValue(), DL, Ty));
   } else {
-    // FIXME: https://github.com/llvm/llvm-project/issues/26431
-    DAG.getContext()->emitError("Can't expand stackmap operand");
+    // FIXME: There are a couple of problems with expanding non-constants for
+    // stackmaps:
+    //  - https://github.com/llvm/llvm-project/issues/26431
+    //  - https://github.com/llvm/llvm-project/issues/55957
+    DAG.getContext()->emitError(
+        "expanding this stackmap operand is unimplemented");
   }
 
   // Copy remaining operands.
diff --git a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
index a93ffac1e1bb..80d95e3b66a5 100644
--- a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
+++ b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
@@ -29,7 +29,7 @@
 ;     CHECK-NEXT:   .long {{.*}}
 ;     CHECK-NEXT:   .short {{.*}}
 ;     NumLocations
-;     CHECK-NEXT:   .short 5
+;     CHECK-NEXT:   .short 7
 ;     Location[NumLocations]
 ;       Location[0]
 ;         CHECK-NEXT: .byte   1
@@ -48,7 +48,7 @@
 ;       Location[2]
 ;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
-;         CHECK-NEXT: .short  16
+;         CHECK-NEXT: .short  1
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
@@ -60,26 +60,51 @@
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
 ;       Location[4]
+;         CHECK-NEXT: .byte   1
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  16
+;         CHECK-NEXT: .short  {{.*}}
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
+;       Location[5]
 ;         CHECK-NEXT: .byte   4
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  8
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   66
+;       Location[4]
+;         CHECK-NEXT: .byte   1
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  4
+;         CHECK-NEXT: .short  {{.*}}
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
+
+@p32 = external global i8 addrspace(270)*
 
 declare void @llvm.experimental.stackmap(i64, i32, ...)
 
 define dso_local i32 @main(i32 %argc, i8** %argv) {
 entry:
-  %intreg = icmp eq i32 %argc, 5
+  %i1reg = icmp eq i32 %argc, 5
+  %i7reg = zext i1 %i1reg to i7
+  %i128reg = zext i1 %i1reg to i128
   %halfreg = sitofp i32 %argc to half
+  %ptr32 = load i8 addrspace(270)*, i8 addrspace(270)** @p32
   call void (i64, i32, ...) @llvm.experimental.stackmap(
     i64 0,
     i32 0,
-    i1 %intreg,
+    i1 %i1reg,
     i7 22,
+    i7 %i7reg,
     half 1.0,
     half %halfreg,
-    i128 66)
+    i128 66,
+    ; FIXME: test non-constant i128 once these are fixed:
+    ;  - https://github.com/llvm/llvm-project/issues/26431
+    ;  - https://github.com/llvm/llvm-project/issues/55957
+    ;i128 %i128reg
+    i8 addrspace(270)* %ptr32)
   ret i32 0
 }

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

84

I've added the missing non-constant i7, and I would have done the same for i128, but that's already broken in llvm before my change and my change doesn't fix it either. That can be another diff to fix that.

Also added a 32-bit address space pointer into the test.

[vext01]

Updated the diff to test structs and added some more comments.

Note that structs and vectors were broken before my change, and remain broken for now.

Diff to last version:

diff --git a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
index 80d95e3b66a5..32242ac1239b 100644
--- a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
+++ b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
@@ -29,7 +29,7 @@
 ;     CHECK-NEXT:   .long {{.*}}
 ;     CHECK-NEXT:   .short {{.*}}
 ;     NumLocations
-;     CHECK-NEXT:   .short 7
+;     CHECK-NEXT:   .short 11
 ;     Location[NumLocations]
 ;       Location[0]
 ;         CHECK-NEXT: .byte   1
@@ -80,18 +80,50 @@
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
+;       Location[5]
+;         CHECK-NEXT: .byte   4
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  8
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
+;       Location[6]
+;         CHECK-NEXT: .byte   1
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  4
+;         CHECK-NEXT: .short  {{.*}}
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
+;       Location[7]
+;         CHECK-NEXT: .byte   4
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  8
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
+;       Location[6]
+;         CHECK-NEXT: .byte   1
+;         CHECK-NEXT: .byte   0
+;         CHECK-NEXT: .short  1
+;         CHECK-NEXT: .short  {{.*}}
+;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .long   0
 
 @p32 = external global i8 addrspace(270)*
 
+%struct1 = type {i32, i64}
+%struct2 = type {i1, i1, i1}
+
 declare void @llvm.experimental.stackmap(i64, i32, ...)
 
 define dso_local i32 @main(i32 %argc, i8** %argv) {
 entry:
   %i1reg = icmp eq i32 %argc, 5
   %i7reg = zext i1 %i1reg to i7
-  %i128reg = zext i1 %i1reg to i128
   %halfreg = sitofp i32 %argc to half
   %ptr32 = load i8 addrspace(270)*, i8 addrspace(270)** @p32
+  %structreg1 = insertvalue %struct1 zeroinitializer, i32 %argc, 0
+  %structreg2 = insertvalue %struct2 zeroinitializer, i1 %i1reg, 0
   call void (i64, i32, ...) @llvm.experimental.stackmap(
     i64 0,
     i32 0,
@@ -101,10 +133,20 @@ entry:
     half 1.0,
     half %halfreg,
     i128 66,
+    ; FIXME: fix and test vectors. At the moment even legally sized vectors
+    ; are broken:
+    ; https://github.com/llvm/llvm-project/issues/55613
+    ;
     ; FIXME: test non-constant i128 once these are fixed:
     ;  - https://github.com/llvm/llvm-project/issues/26431
     ;  - https://github.com/llvm/llvm-project/issues/55957
-    ;i128 %i128reg
-    i8 addrspace(270)* %ptr32)
+    i8 addrspace(270)* %ptr32,
+    ; FIXME: The stackmap record generated for structs is incorrect:
+    ;  - https://github.com/llvm/llvm-project/issues/55649
+    ;  - https://github.com/llvm/llvm-project/issues/55957
+    %struct1 zeroinitializer,
+    %struct1 %structreg1,
+    %struct2 zeroinitializer,
+    %struct2 %structreg2)
   ret i32 0
 }

Let's move this forward. Can we put this into main?

[arsenm]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

2314

I'm a bit surprised this can unconditionally zero extend. I'd expect ANY_EXTEND or to have to consider some ABI property

5510–5517

I'd rather not add a special case to make this work only for constants

5522–5523

My skimming of the issue says this is an issue with the lowering from IR, so why do you need to error here? In any case I think emitError should be reserved for cases where reasonable user code ran into a truly unhandleable case. It would be better to just let it fall through and hit the ordinary legalize fatal error

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2274–2277

You're still fixing this into TargetConstant at selection time instead of upfront when lowering from the IR

[vext01]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

2314

Perhaps not some ABI property, as a call to llvm.experimental.stackmap isn't really a call, so there is no ABI to speak of.

I'll look into ANY_EXTEND. I'm not familiar.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2274–2277

This code is selecting the stackmap live variables.

The (non-frameindex) live variables can't be emitted to target constants at DAG-build time or they won't get legalized and that's the problem that this change is trying to address.

Correct me if I'm wrong though.

[vext01]

OK, on IRC @arsenm and I agreed that for the cases that we can't yet handle, we should see the same legaliser error that we see in-tree now.o

This change implements that, albeit slightly clumsily with a goto. If you simply don't expand the problematic operand, as I first tried, the legaliser gets stuck in an infinite loop. Note also that goto default is not valid in C/C++, so had to add another label.

We do still have to special-case constants I'm afraid.

Also fixed the EXTEND_ANY thing and fixed some incorrectly numbered comments in the test.

Diff to previous:

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index 7ed2808600cf..44e6a9e5a6a6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -2311,7 +2311,10 @@ SDValue DAGTypeLegalizer::PromoteIntOp_SET_ROUNDING(SDNode *N) {
 SDValue DAGTypeLegalizer::PromoteIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
   assert(OpNo > 1); // Because the first two arguments are guaranteed legal.
   SmallVector<SDValue> NewOps(N->ops().begin(), N->ops().end());
-  NewOps[OpNo] = ZExtPromotedInteger(N->getOperand(OpNo));
+  SDValue Operand = N->getOperand(OpNo);
+  EVT NVT = TLI.getTypeToTransformTo(*DAG.getContext(), Operand.getValueType());
+  NewOps[OpNo] =
+      DAG.getNode(ISD::ANY_EXTEND, SDLoc(N), NVT, N->getOperand(OpNo));
   return SDValue(DAG.UpdateNodeOperands(N, NewOps), 0);
 }
 
@@ -4631,6 +4634,7 @@ bool DAGTypeLegalizer::ExpandIntegerOperand(SDNode *N, unsigned OpNo) {
     return false;
 
   switch (N->getOpcode()) {
+  fail:
   default:
   #ifndef NDEBUG
     dbgs() << "ExpandIntegerOperand Op #" << OpNo << ": ";
@@ -4665,7 +4669,11 @@ bool DAGTypeLegalizer::ExpandIntegerOperand(SDNode *N, unsigned OpNo) {
 
   case ISD::ATOMIC_STORE:      Res = ExpandIntOp_ATOMIC_STORE(N); break;
   case ISD::STACKMAP:
-    Res = ExpandIntOp_STACKMAP(N, OpNo);
+    Optional<SDValue> MaybeRes = ExpandIntOp_STACKMAP(N, OpNo);
+    if (MaybeRes.hasValue())
+      Res = MaybeRes.getValue();
+    else
+      goto fail;
     break;
   }
 
@@ -5496,7 +5504,8 @@ SDValue DAGTypeLegalizer::PromoteIntOp_CONCAT_VECTORS(SDNode *N) {
   return DAG.getBuildVector(N->getValueType(0), dl, NewOps);
 }
 
-SDValue DAGTypeLegalizer::ExpandIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
+Optional<SDValue> DAGTypeLegalizer::ExpandIntOp_STACKMAP(SDNode *N,
+                                                         unsigned OpNo) {
   assert(OpNo > 1);
 
   SDValue Op = N->getOperand(OpNo);
@@ -5508,19 +5517,21 @@ SDValue DAGTypeLegalizer::ExpandIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
     NewOps.push_back(N->getOperand(I));
 
   if (Op->getOpcode() == ISD::Constant) {
-    // FIXME: https://github.com/llvm/llvm-project/issues/55609
     ConstantSDNode *CN = cast<ConstantSDNode>(Op);
     EVT Ty = Op.getValueType();
-    NewOps.push_back(
-        DAG.getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
-    NewOps.push_back(DAG.getTargetConstant(CN->getZExtValue(), DL, Ty));
+    if (CN->getConstantIntValue()->getValue().getActiveBits() < 64) {
+      NewOps.push_back(
+          DAG.getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
+      NewOps.push_back(DAG.getTargetConstant(CN->getZExtValue(), DL, Ty));
+    } else {
+      // FIXME: https://github.com/llvm/llvm-project/issues/55609
+      return Optional<SDValue>();
+    }
   } else {
-    // FIXME: There are a couple of problems with expanding non-constants for
-    // stackmaps:
+    // FIXME: Non-constant operands are not yet handled:
     //  - https://github.com/llvm/llvm-project/issues/26431
     //  - https://github.com/llvm/llvm-project/issues/55957
-    DAG.getContext()->emitError(
-        "expanding this stackmap operand is unimplemented");
+    return Optional<SDValue>();
   }
 
   // Copy remaining operands.
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 2807b7f5ae68..bc8257c9ec4c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -494,7 +494,7 @@ private:
   SDValue ExpandIntOp_RETURNADDR(SDNode *N);
   SDValue ExpandIntOp_ATOMIC_STORE(SDNode *N);
   SDValue ExpandIntOp_SPLAT_VECTOR(SDNode *N);
-  SDValue ExpandIntOp_STACKMAP(SDNode *N, unsigned OpNo);
+  Optional<SDValue> ExpandIntOp_STACKMAP(SDNode *N, unsigned OpNo);
 
   void IntegerExpandSetCCOperands(SDValue &NewLHS, SDValue &NewRHS,
                                   ISD::CondCode &CCCode, const SDLoc &dl);
diff --git a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
index 32242ac1239b..bc624be5318e 100644
--- a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
+++ b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
@@ -73,35 +73,35 @@
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   66
-;       Location[4]
+;       Location[6]
 ;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  4
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
-;       Location[5]
+;       Location[7]
 ;         CHECK-NEXT: .byte   4
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  8
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
-;       Location[6]
+;       Location[8]
 ;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  4
 ;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
-;       Location[7]
+;       Location[9]
 ;         CHECK-NEXT: .byte   4
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  8
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
-;       Location[6]
+;       Location[10]
 ;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
 ;         CHECK-NEXT: .short  1

LGTY?

[arsenm]

In D125680#3615526, @vext01 wrote:

Diff to previous:

FYI Phabricator gives you this if you look at the history tab

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2274–2277

This is a pretty weird behavior, but does the type actually matter? Could you just unconditionally use 64-bit constants?

[vext01]

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

2274–2277

Sorry for the delay. I was away at the end of last week.

I've just tried to implement your suggestion, it has two undesirable effects:

• It changes the way large constant structs which get split into e.g. mergevalues(constant, constant) are codegenned. It appears that the constituent parts of the struct get allocated to registers and not constants.
• For FastISel, some some small constants get reported as long constants in the stackmap record.
• For FastISel It causes large constants to be emitted in a different order than in SelectionDAGISel, which means that the test files which test both backends cannot succeed.

This is too much scary breakage, so I suggest we emit constants as we did before.

Here is the diff (to previous) that I was working on before I abandoned it:

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index b9816bb34bc9..497c4c77880a 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -9449,6 +9449,10 @@ void SelectionDAGBuilder::visitStackmap(const CallInst &CI) {
       const TargetLowering &TLI = DAG.getTargetLoweringInfo();
       Ops.push_back(DAG.getTargetFrameIndex(
           FI->getIndex(), TLI.getFrameIndexTy(DAG.getDataLayout())));
+    } else if (ConstantSDNode *C = dyn_cast<ConstantSDNode>(Op)) {
+		// Directly emit a 64-bit target constant.
+		Ops.push_back(DAG.getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
+		Ops.push_back(DAG.getTargetConstant(C->getZExtValue(), DL, MVT::i64));
     } else {
       // Otherwise emit a target independent node to be legalised.
       Ops.push_back(getValue(CI.getArgOperand(I)));
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
index aff4e9a94fef..936f48f34870 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
@@ -2269,15 +2269,7 @@ void SelectionDAGISel::Select_STACKMAP(SDNode *N) {
     // FrameIndex nodes should have been directly emitted to TargetFrameIndex
     // nodes at DAG-construction time.
     assert(OpNode->getOpcode() != ISD::FrameIndex);
-
-    if (OpNode->getOpcode() == ISD::Constant) {
-      Ops.push_back(
-          CurDAG->getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
-      O = CurDAG->getTargetConstant(
-          cast<ConstantSDNode>(OpNode)->getZExtValue(), DL, It->getValueType());
-    } else {
-      O = *It;
-    }
+    O = *It;
     Ops.push_back(O);
   }
 
diff --git a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
index bc624be5318e..a2d629fce00c 100644
--- a/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
+++ b/llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
@@ -81,10 +81,10 @@
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
 ;       Location[7]
-;         CHECK-NEXT: .byte   4
+;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
-;         CHECK-NEXT: .short  8
-;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .short  4
+;         CHECK-NEXT: .short  {{.*}}
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
 ;       Location[8]
@@ -95,10 +95,10 @@
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
 ;       Location[9]
-;         CHECK-NEXT: .byte   4
+;         CHECK-NEXT: .byte   1
 ;         CHECK-NEXT: .byte   0
-;         CHECK-NEXT: .short  8
-;         CHECK-NEXT: .short  0
+;         CHECK-NEXT: .short  1
+;         CHECK-NEXT: .short  2
 ;         CHECK-NEXT: .short  0
 ;         CHECK-NEXT: .long   0
 ;       Location[10]
diff --git a/llvm/test/CodeGen/X86/stackmap-fast-isel.ll b/llvm/test/CodeGen/X86/stackmap-fast-isel.ll
index dd25065f3063..558029cdf1c5 100644
--- a/llvm/test/CodeGen/X86/stackmap-fast-isel.ll
+++ b/llvm/test/CodeGen/X86/stackmap-fast-isel.ll
@@ -29,8 +29,8 @@
 ; CHECK-NEXT:   .quad 4
 
 ; Large Constants
-; CHECK-NEXT:   .quad   2147483648
 ; CHECK-NEXT:   .quad   4294967295
+; CHECK-NEXT:   .quad   2147483648
 ; CHECK-NEXT:   .quad   4294967296
 
 ; Callsites
@@ -46,14 +46,14 @@
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   -1
+; CHECK-NEXT:   .long   65535
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   0
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   -1
+; CHECK-NEXT:   .long   65535
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   0
@@ -76,19 +76,19 @@
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .long   2147483647
 ; SmallConstant
-; CHECK-NEXT:   .byte   4
+; CHECK-NEXT:   .byte   5
 ; CHECK-NEXT:   .byte   0
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   -1
+; CHECK-NEXT:   .long   0
 ; SmallConstant
-; CHECK-NEXT:   .byte   4
+; CHECK-NEXT:   .byte   5
 ; CHECK-NEXT:   .byte   0
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   -1
+; CHECK-NEXT:   .long   0
 ; SmallConstant
 ; CHECK-NEXT:   .byte   4
 ; CHECK-NEXT:   .byte   0
@@ -102,14 +102,14 @@
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   0
+; CHECK-NEXT:   .long   1
 ; LargeConstant at index 1
 ; CHECK-NEXT:   .byte   5
 ; CHECK-NEXT:   .byte   0
 ; CHECK-NEXT:   .short  8
 ; CHECK-NEXT:   .short  0
 ; CHECK-NEXT:   .short  0
-; CHECK-NEXT:   .long   1
+; CHECK-NEXT:   .long   0
 ; LargeConstant at index 2
 ; CHECK-NEXT:   .byte   5
 ; CHECK-NEXT:   .byte   0

[arsenm]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

4675–4676

Don't see why you are introducing this goto

5507

Don't see why you're using Optional. Just use SDValue() like everything else?

[vext01]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

4675–4676

It's because it's the only path to the error state contained in the default branch of the switch.

In an earlier conversation we said that we wanted to trigger the default "can't legalise this" error if we encounter a case we cannot handle.

5507

If we return SDValue() then we are signalling that the handler has already replaced the node with a new node with the operand in question legalised.

If we do that, then the legaliser will get stuck in an infinite loop trying to legalise the operand that we cannot handle, since we never actually replace the operand.

I introduced the optional so that we could signal the error state and break out of the legaliser loop (i.e. the goto you questioned).

[arsenm]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

4675–4676

But there's no actual failure case here so I don't understand what this is doing. This is always producing something

5527–5528

Just leave the node as-is? It will fail to legalize later

[vext01]

Removed clunky error handling.

[vext01]

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

5527–5528

Sorry, there was some confusion at my end.

Leaving the node as-is (i.e. return SDValue(N, 0);) is the case that will cause an infinite loop in the legaliser. Returning a "null" value (i.e. return SDValue()) is fine though.

If I put a print at the top of this function:

errs() << __func__ << ": " << OpNo << ": "; N->dump();

And update the return SDValue() to return SDValue(N, 0);, and then use an input that covers one of those cases, e.g.:

declare void @llvm.experimental.stackmap(i64, i32, ...)

define void @f() {
  call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 0, i128 9223372036854775808)
  ret void
}

Then we see the function called indefinitely, trying to legalise the same operand forever:

ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
ExpandIntOp_STACKMAP: 4: t8: ch,glue = stackmap t3, t3:1, TargetConstant:i64<0>, TargetConstant:i32<0>, Constant:i128<9223372036854775808>
...

This is due to the special casing in DAGTypeLegalizer::ExpandIntegerOperand():

Res = ExpandIntOp_STACKMAP(N, OpNo);
...
    // If the result is null, the sub-method took care of registering results etc.                                                                                                                                                            
    if (!Res.getNode()) return false;                                                                                                                                                                                                         
                                                                                                                                                                                                                                              
    // If the result is N, the sub-method updated N in place.  Tell the legalizer                                                                                                                                                             
    // core about this.                                                                                                                                                                                                                       
    if (Res.getNode() == N)                                                                                                                                                                                                                   
      return true;

Returning SDValue() hits the return false, whereas return SDValue(N, 0) hits return true.

The return value is whether or not the operand needs another round of legalisation. At the call site of ExpandIntegerOperand(), there is:

NeedsReanalyzing = ExpandIntegerOperand(N, i);

So we need to return false to:

• not loop forever
• bail out with the right error.

So I believe return SDValue() to be correct. The latest diff reflects this.

[vext01]

Small tweak.

[vext01]

This landed in ed8ef65

llvm/include/llvm/CodeGen/ISDOpcodes.h

 #define BEGIN_REGISTER_VP_SDNODE(VPSDID, ...) VPSDID,
 #include "llvm/IR/VPIntrinsics.def"
 
+  // The `llvm.experimental.stackmap` intrinsic.
+  // Operands: input chain, glue, <id>, <numShadowBytes>, [live0[, live1...]]
+  // Outputs: output chain, glue
+  STACKMAP,
+
   /// BUILTIN_OP_END - This must be the last enum value in this list.
   /// The target-specific pre-isel opcode values start here.
   BUILTIN_OP_END

llvm/include/llvm/CodeGen/SelectionDAGISel.h

 
   void Select_FREEZE(SDNode *N);
   void Select_ARITH_FENCE(SDNode *N);
+  void Select_STACKMAP(SDNode *N);
 
 private:
   void DoInstructionSelection();

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

     break;
 
   case ISD::SET_ROUNDING: Res = PromoteIntOp_SET_ROUNDING(N); break;
+  case ISD::STACKMAP:
+    Res = PromoteIntOp_STACKMAP(N, OpNo);
+    break;
   }
 
   // If the result is null, the sub-method took care of registering results etc.

   return SDValue(DAG.UpdateNodeOperands(N, N->getOperand(0), Op), 0);
 }
 
+SDValue DAGTypeLegalizer::PromoteIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
+  assert(OpNo > 1); // Because the first two arguments are guaranteed legal.
+  SmallVector<SDValue> NewOps(N->ops().begin(), N->ops().end());
+  NewOps[OpNo] = ZExtPromotedInteger(N->getOperand(OpNo));
+  return SDValue(DAG.UpdateNodeOperands(N, NewOps), 0);
+}
+
 //===----------------------------------------------------------------------===//
 //  Integer Result Expansion
 //===----------------------------------------------------------------------===//

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

   SDValue PromoteIntOp_VECREDUCE(SDNode *N);
   SDValue PromoteIntOp_VP_REDUCE(SDNode *N, unsigned OpNo);
   SDValue PromoteIntOp_SET_ROUNDING(SDNode *N);
+  SDValue PromoteIntOp_STACKMAP(SDNode *N, unsigned OpNo);
 
   void PromoteSetCCOperands(SDValue &LHS,SDValue &RHS, ISD::CondCode Code);
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

   }
 }
 
-/// Lower llvm.experimental.stackmap directly to its target opcode.
+/// Lower llvm.experimental.stackmap.
 void SelectionDAGBuilder::visitStackmap(const CallInst &CI) {
-  // void @llvm.experimental.stackmap(i32 <id>, i32 <numShadowBytes>,
+  // void @llvm.experimental.stackmap(i64 <id>, i32 <numShadowBytes>,
   //                                  [live variables...])
 
   assert(CI.getType()->isVoidTy() && "Stackmap cannot return a value.");

   Chain = DAG.getCALLSEQ_START(getRoot(), 0, 0, DL);
   InFlag = Chain.getValue(1);
 
-  // Add the <id> and <numBytes> constants.
-  SDValue IDVal = getValue(CI.getOperand(PatchPointOpers::IDPos));
-  Ops.push_back(DAG.getTargetConstant(
-                  cast<ConstantSDNode>(IDVal)->getZExtValue(), DL, MVT::i64));
-  SDValue NBytesVal = getValue(CI.getOperand(PatchPointOpers::NBytesPos));
-  Ops.push_back(DAG.getTargetConstant(
-                  cast<ConstantSDNode>(NBytesVal)->getZExtValue(), DL,
-                  MVT::i32));
-
-  // Push live variables for the stack map.
-  addStackMapLiveVars(CI, 2, DL, Ops, *this);
-
-  // We are not pushing any register mask info here on the operands list,
-  // because the stackmap doesn't clobber anything.
+  // Add the STACKMAP operands, starting with DAG house-keeping.
+  std::vector<SDValue> Operands;
+  Operands.push_back(Chain);
+  Operands.push_back(InFlag);
 
-  // Push the chain and the glue flag.
-  Ops.push_back(Chain);
-  Ops.push_back(InFlag);
+  // Add the <id>, <numShadowBytes> and live variables arguments.
+  for (unsigned I = 0; I < CI.arg_size(); I++)
+    Operands.push_back(getValue(CI.getArgOperand(I)));
 
   // Create the STACKMAP node.
   SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
-  SDNode *SM = DAG.getMachineNode(TargetOpcode::STACKMAP, DL, NodeTys, Ops);
-  Chain = SDValue(SM, 0);
+  Chain = DAG.getNode(ISD::STACKMAP, DL, NodeTys, Operands);
   InFlag = Chain.getValue(1);
 
   Chain = DAG.getCALLSEQ_END(Chain, NullPtr, NullPtr, InFlag, DL);

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

   case ISD::SDID:                                                              \
     return #NAME;
 #include "llvm/IR/VPIntrinsics.def"
+
+  case ISD::STACKMAP:
+    return "stackmap";
   }
 }
 

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

 #include "llvm/CodeGen/SchedulerRegistry.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
+#include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/SwiftErrorValueTracking.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"

                        N->getOperand(0));
 }
 
+void SelectionDAGISel::Select_STACKMAP(SDNode *N) {
+  std::vector<SDValue> Ops;
+  auto *It = N->op_begin();
+  SDLoc DL(N);
+
+  // Stash the chain and glue operands so we can move them to the end.
+  SDValue Chain = *It++;
+  SDValue InFlag = *It++;
+
+  // <id> operand.
+  SDValue ID = It->get();
+  assert(ID.getOpcode() == ISD::Constant);
+  assert(ID.getValueType() == MVT::i64);
+  SDValue IDConst = CurDAG->getTargetConstant(
+      cast<ConstantSDNode>(ID)->getZExtValue(), DL, ID.getValueType());
+  Ops.push_back(IDConst);
+  It++;
+
+  // <numBytes> operand.
+  SDValue Shad = It->get();
+  assert(Shad.getOpcode() == ISD::Constant);
+  assert(Shad.getValueType() == MVT::i32);
+  SDValue ShadConst = CurDAG->getTargetConstant(
+      cast<ConstantSDNode>(Shad)->getZExtValue(), DL, Shad.getValueType());
+  Ops.push_back(ShadConst);
+  It++;
+
+  // Live variable operands.
+  for (; It != N->op_end(); It++) {
+    SDNode *OpNode = It->getNode();
+    SDValue O;
+    if (OpNode->getOpcode() == ISD::Constant) {
+      Ops.push_back(
+          CurDAG->getTargetConstant(StackMaps::ConstantOp, DL, MVT::i64));
+      O = CurDAG->getTargetConstant(
+          cast<ConstantSDNode>(OpNode)->getZExtValue(), DL, It->getValueType());
+    } else if (OpNode->getOpcode() == ISD::FrameIndex) {
+      FrameIndexSDNode *FI = cast<FrameIndexSDNode>(OpNode);
+      O = CurDAG->getTargetFrameIndex(FI->getIndex(), It->getValueType());
+    } else {
+      // Otherwise it's a register.
+      // XXX ^^ is this true?
+      O = *It; // XXX Also does this guarantee register selection?
+    }
+    Ops.push_back(O);
+  }
+
+  Ops.push_back(Chain);
+  Ops.push_back(InFlag);
+
+  SDVTList NodeTys = CurDAG->getVTList(MVT::Other, MVT::Glue);
+  CurDAG->SelectNodeTo(N, TargetOpcode::STACKMAP, NodeTys, Ops);
+}
+
 /// GetVBR - decode a vbr encoding whose top bit is set.
 LLVM_ATTRIBUTE_ALWAYS_INLINE static uint64_t
 GetVBR(uint64_t Val, const unsigned char *MatcherTable, unsigned &Idx) {

   case ISD::ARITH_FENCE:
     Select_ARITH_FENCE(NodeToMatch);
     return;
+  case ISD::STACKMAP:
+    Select_STACKMAP(NodeToMatch);
+    return;
   }
 
   assert(!NodeToMatch->isMachineOpcode() && "Node already selected!");

llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll

+; RUN: llc -debug-only=legalize-types -print-after=finalize-isel %s -O1 -mtriple=x86_64-unknown-unknown -o /dev/null 2>&1 | FileCheck %s
+
+
+declare void @llvm.experimental.stackmap(i64, i32, ...)
+
+define dso_local i32 @main(i32 %argc, i8** %argv) {
+entry:
+  %x = icmp eq i32 %argc, 5
+  call void (i64, i32, ...) @llvm.experimental.stackmap(i64 0, i32 0, i1 %x, i7 2)
+  ; CHECK: Legalizing node: t14: ch,glue = stackmap t10, t10:1, Constant:i64<0>, Constant:i32<0>, t7, Constant:i7<2>
+  ; CHECK-NEXT: Analyzing result type: ch
+  ; CHECK-NEXT: Legal result type
+  ; CHECK-NEXT: Analyzing result type: glue
+  ; CHECK-NEXT: Legal result type
+  ; CHECK-NEXT: Analyzing operand: t10: ch,glue = callseq_start t0, TargetConstant:i64<0>, TargetConstant:i64<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t10: ch,glue = callseq_start t0, TargetConstant:i64<0>, TargetConstant:i64<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t11: i64 = Constant<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t12: i32 = Constant<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t7: i1 = setcc t2, Constant:i32<5>, seteq:ch
+  ; CHECK-NEXT: Promote integer operand: t14: ch,glue = stackmap t10, t10:1, Constant:i64<0>, Constant:i32<0>, t7, Constant:i7<2>
+
+  ; CHECK: Legalizing node: t14: ch,glue = stackmap t10, t10:1, Constant:i64<0>, Constant:i32<0>, t23, Constant:i7<2>
+  ; CHECK-NEXT: Analyzing result type: ch
+  ; CHECK-NEXT: Legal result type
+  ; CHECK-NEXT: Analyzing result type: glue
+  ; CHECK-NEXT: Legal result type
+  ; CHECK-NEXT: Analyzing operand: t10: ch,glue = callseq_start t0, TargetConstant:i64<0>, TargetConstant:i64<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t10: ch,glue = callseq_start t0, TargetConstant:i64<0>, TargetConstant:i64<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t11: i64 = Constant<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t12: i32 = Constant<0>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t23: i8 = and t21, Constant:i8<1>
+  ; CHECK-NEXT: Legal operand
+  ; CHECK-NEXT: Analyzing operand: t13: i7 = Constant<2>
+  ; CHECK-NEXT: Promote integer operand: t14: ch,glue = stackmap t10, t10:1, Constant:i64<0>, Constant:i32<0>, t23, Constant:i7<2>
+
+  ; CHECK: # Machine code for function main
+  ; CHECK: bb.0.entry:
+  ; CHECK: STACKMAP 0, 0, killed %3:gr8, 2, 2, implicit-def dead early-clobber $r11
+  ; CHECK: # End machine code for function main
+  ret i32 0
+}