[llvm] [DAGCombiner] Restrict mergeTruncStores to pre-type-legalization combines (PR #222245)

[Chung-Yi Chen via llvm-commits]

================
@@ -9946,10 +9946,14 @@ static SDValue stripTruncAndExt(SDValue Value) {
 SDValue DAGCombiner::mergeTruncStores(StoreSDNode *N) {
   // The matching looks for "store (trunc x)" patterns that appear early but are
   // likely to be replaced by truncating store nodes during combining.
+  // This must be restricted to the combine rounds that run before type
+  // legalization: the fold creates a truncate to the wide type plus an optional
+  // bswap/rotate of it, and those are only allowed to have an illegal type
+  // while types have not been legalized yet.
----------------
ADNRs wrote:

After some investigation, I think that I somehow misunderstood the meaning of the DAG node, so the comment seems weird. But the patch still works, and a better fix will be discussed later. Sorry for my shallow understanding of the backend :(

The problem is not about bswap, but about the type `WideVT`. The code snippet below is where the problematic bswap is created. The key is that an i16 node is created by `mergeTruncStores()` after type legalization, but i16 (`WideVT`) is not legal on AArch64.

```C++
if (NeedBswap) {
    SourceValue = DAG.getNode(ISD::BSWAP, DL, WideVT, SourceValue);
  } else if (NeedRotate) {
```

The original `LegalOperations` guard was added by 1d0fa798248f to restrict this fold to early combining, but `LegalOperations` is only set from the `AfterLegalizeVectorOps` round onwards, so the fold still runs in the `AfterLegalizeTypes` round. That round already requires every node to have a legal type, and that is where the i16 bswap gets created. Guarding on `LegalTypes` makes the restriction match its description and prevents the illegal node.

Now I think a better fix is to guard the bswap creation instead of moving from `LegalOperations` to `LegalTypes`:

```C++
if (NeedBswap) {
    if (!isTypeLegal(WideVT))
        return SDValue();
    SourceValue = DAG.getNode(ISD::BSWAP, DL, WideVT, SourceValue);
  } else if (NeedRotate) {
```

I am not sure whether rotate will suffer a similar bug. I cannot craft an example like the test case that crashes on rotate. Should we guard the rotate creation, too?

What do you think? If you also prefer this new fix, I'll proceed. If not, I'll rewrite the comment you quoted.

https://github.com/llvm/llvm-project/pull/222245