This is an archive of the discontinued LLVM Phabricator instance.

Correctly legalise stackmap operands
ClosedPublic

Authored by vext01 on May 16 2022, 4:25 AM.

Details

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)

Diff Detail

Event Timeline

vext01 created this revision.May 16 2022, 4:25 AM
Herald added a project: Restricted Project. · View Herald TranscriptMay 16 2022, 4:25 AM
vext01 requested review of this revision.May 16 2022, 4:25 AM
Herald added a project: Restricted Project. · View Herald TranscriptMay 16 2022, 4:25 AM

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 added a subscriber: reames.

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 added inline comments.May 16 2022, 6:04 AM
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
3

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

10

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

vext01 updated this revision to Diff 429695.May 16 2022, 6:05 AM

Updated diff with full context.

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 added inline comments.May 16 2022, 6:20 AM
llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
2279–2280

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

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
3

At what level should I check?

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

10

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

arsenm added inline comments.May 16 2022, 6:44 AM
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
3

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

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 added inline comments.May 16 2022, 7:10 AM
llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
3

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 added inline comments.May 16 2022, 2:03 PM
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
3

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

vext01 added inline comments.May 17 2022, 2:00 AM
llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
2286

And is that correct?

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

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 added inline comments.May 17 2022, 2:02 AM
llvm/test/CodeGen/X86/selectiondag-stackmap-legalize.ll
3

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

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 added inline comments.May 17 2022, 7:30 AM
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

Hi Philip,

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...

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 updated this revision to Diff 430370.EditedMay 18 2022, 7:06 AM

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.
  • 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 updated this revision to Diff 430371.May 18 2022, 7:08 AM

Sorry, bodged the last diff upload.

vext01 added a comment.EditedMay 18 2022, 7:17 AM

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?

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.

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 updated this revision to Diff 435111.Jun 8 2022, 5:30 AM

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:

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 updated this revision to Diff 437863.Jun 17 2022, 5:59 AM

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)

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

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

4671

Nit: and here too

4686

Nit: extra empty line

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
9422

Use existing Ops instead of introducing new vector?

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
490

Single line again

vext01 added inline comments.Jun 21 2022, 1:39 AM
llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp
2920

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

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 updated this revision to Diff 439000.Jun 22 2022, 6:33 AM

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

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

arsenm added inline comments.Jun 23 2022, 6:12 AM
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 added inline comments.Jun 23 2022, 6:34 AM
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 updated this revision to Diff 439406.Jun 23 2022, 8:16 AM

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 this revision to Diff 439691.Jun 24 2022, 3:32 AM

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 added inline comments.Jun 27 2022, 6:12 AM
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

5512–5519

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

5524–5525

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 added inline comments.Jun 27 2022, 7:04 AM
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 updated this revision to Diff 440603.Jun 28 2022, 7:03 AM

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?

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 added inline comments.Jul 4 2022, 7:00 AM
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 added inline comments.Jul 5 2022, 8:34 AM
llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
4673–4674

Don't see why you are introducing this goto

5501

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

vext01 added inline comments.Jul 5 2022, 9:14 AM
llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
4673–4674

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.

5501

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 added inline comments.Jul 5 2022, 10:15 AM
llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
4673–4674

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

5521–5522

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

vext01 updated this revision to Diff 442478.Jul 6 2022, 2:15 AM

Removed clunky error handling.

vext01 added inline comments.Jul 6 2022, 2:52 AM
llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
5521–5522

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 updated this revision to Diff 442540.Jul 6 2022, 6:03 AM

Small tweak.

arsenm accepted this revision.Jul 6 2022, 6:44 AM
This revision is now accepted and ready to land.Jul 6 2022, 6:44 AM
vext01 closed this revision.Jul 18 2022, 8:27 AM

This landed in ed8ef65