[CodeGen] Fix type legalization crash on vectors of reference type pointers

[hoodmane]

Fixes #69894

When the SLP vectorizer cost-models consecutive externref loads/stores, it forms an vector and asks the target to legalize it. Type legalization recovers the vector element type with DataLayout-free EVT helpers which route through EVT::getEVT()/MVT::getVT(). Those cannot lower a pointer to an MVT without a DataLayout and hit llvm_unreachable("Unknown type!")

EVT::getTypeForEVT maps MVT::externref / MVT::funcref to opaque pointers in address spaces 10 and 20, so we can fix the crash by making that mapping invertible in MVT::getVT. The mapping is only applied on the HandleUnknown=false path used by that round-trip; HandleUnknown=true callers keep getting MVT::Other, so behavior for every other target is unchanged.

Address spaces 10/20 are not reserved for WebAssembly, but a target only reaches this code with such a pointer after producing an externref/funcref MVT, which only WebAssembly does.

With the type legalizable, scalarizes, so the vectorizer leaves the reference types scalar rather than crashing.

It would be much nicer if this could be gated on the WebAssembly target specifically but as far as I can tell there is no way to ask about the target in getVT().

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-backend-webassembly

Author: Hood Chatham (hoodmane)

Changes

When the SLP vectorizer cost-models consecutive externref loads/stores, it forms an <N x externref> vector and asks the target to legalize it. Type legalization recovers the vector element type with DataLayout-free EVT helpers which route through EVT::getEVT()/MVT::getVT(). Those cannot lower a pointer to an MVT without a DataLayout and hit llvm_unreachable("Unknown type!")

EVT::getTypeForEVT maps MVT::externref / MVT::funcref to opaque pointers in address spaces 10 and 20, so we can fix the crash by making that mapping invertible in MVT::getVT. The mapping is only applied on the HandleUnknown=false path used by that round-trip; HandleUnknown=true callers keep getting MVT::Other, so behavior for every other target is unchanged.

Address spaces 10/20 are not reserved for WebAssembly, but a target only reaches this code with such a pointer after producing an externref/funcref MVT, which only WebAssembly does.

With the type legalizable, <N x externref> scalarizes, so the vectorizer leaves the reference types scalar rather than crashing.

It would be much nicer if this could be gated on the WebAssembly target specifically but as far as I can tell there is no way to ask about the target in getVT().

---

Full diff: https://github.com/llvm/llvm-project/pull/202810.diff

3 Files Affected:

• (modified) llvm/lib/CodeGen/ValueTypes.cpp (+22)
• (added) llvm/test/CodeGen/WebAssembly/externref-phi.ll (+57)
• (added) llvm/test/Transforms/SLPVectorizer/WebAssembly/externref-no-vectorize.ll (+48)

diff --git a/llvm/lib/CodeGen/ValueTypes.cpp b/llvm/lib/CodeGen/ValueTypes.cpp
index e74068e22f4cd..7045219a684b6 100644
--- a/llvm/lib/CodeGen/ValueTypes.cpp
+++ b/llvm/lib/CodeGen/ValueTypes.cpp
@@ -252,6 +252,28 @@ MVT MVT::getVT(Type *Ty, bool HandleUnknown){
   default:
     if (HandleUnknown) return MVT(MVT::Other);
     llvm_unreachable("Unknown type!");
+  case Type::PointerTyID: {
+    if (HandleUnknown) return MVT(MVT::Other);
+    // A pointer normally cannot be lowered to a value type without a
+    // DataLayout, so it is an "unknown" type here. The exception is the
+    // WebAssembly reference types: EVT::getTypeForEVT maps MVT::externref /
+    // MVT::funcref to opaque pointers in address spaces 10 and 20, and that
+    // mapping has to be invertible so that DataLayout-free EVT helpers can
+    // recover the element type of an <N x externref> / <N x funcref> vector
+    // during type legalization instead of crashing.
+    //
+    // Address spaces 10 and 20 are not reserved for WebAssembly, but a target
+    // only reaches this code with such a pointer if it first produced an
+    // externref/funcref MVT, which only WebAssembly does.
+      switch (cast<PointerType>(Ty)->getAddressSpace()) {
+      case 10:
+        return MVT(MVT::externref);
+      case 20:
+        return MVT(MVT::funcref);
+      default:
+        llvm_unreachable("Unknown type!");
+      }
+    }
   case Type::VoidTyID:
     return MVT::isVoid;
   case Type::ByteTyID:
diff --git a/llvm/test/CodeGen/WebAssembly/externref-phi.ll b/llvm/test/CodeGen/WebAssembly/externref-phi.ll
new file mode 100644
index 0000000000000..45010459fd97f
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/externref-phi.ll
@@ -0,0 +1,57 @@
+; RUN: llc < %s --mtriple=wasm32-unknown-unknown -asm-verbose=false -mattr=+reference-types | FileCheck %s
+; RUN: llc < %s --mtriple=wasm64-unknown-unknown -asm-verbose=false -mattr=+reference-types | FileCheck %s
+
+; externref values flow through function arguments, a function return value, and
+; phis that merge them across control flow. They cannot live in linear memory,
+; so they are carried in wasm locals of externref type; a phi is resolved by
+; reassigning the local on the incoming edge.
+;
+; Derived the following C code:
+;
+;   __externref_t foo(void);
+;   void bar(__externref_t);
+;   void test(int flag, __externref_t ref1, __externref_t ref2) {
+;     if (flag) {
+;       ref1 = foo();
+;       ref2 = foo();
+;     }
+;     bar(ref1); bar(ref2);
+;   }
+%externref = type ptr addrspace(10) ;; addrspace 10 is nonintegral
+
+declare %externref @foo()
+declare void @bar(%externref)
+
+; CHECK-LABEL: test:
+; CHECK-NEXT:  .functype       test (i32, externref, externref) -> ()
+; CHECK-NEXT:  block
+; CHECK-NEXT:  local.get       0
+; CHECK-NEXT:  i32.eqz
+; CHECK-NEXT:  br_if           0
+; CHECK-NEXT:  call            foo
+; CHECK-NEXT:  local.set       1
+; CHECK-NEXT:  call            foo
+; CHECK-NEXT:  local.set       2
+; CHECK:       end_block
+; CHECK-NEXT:  local.get       1
+; CHECK-NEXT:  call            bar
+; CHECK-NEXT:  local.get       2
+; CHECK-NEXT:  call            bar
+; CHECK-NEXT:  end_function
+define void @test(i32 %flag, %externref %ref1, %externref %ref2) {
+entry:
+  %c = icmp eq i32 %flag, 0
+  br i1 %c, label %join, label %then
+
+then:
+  %a = tail call %externref @foo()
+  %b = tail call %externref @foo()
+  br label %join
+
+join:
+  %p1 = phi %externref [ %a, %then ], [ %ref1, %entry ]
+  %p2 = phi %externref [ %b, %then ], [ %ref2, %entry ]
+  tail call void @bar(%externref %p1)
+  tail call void @bar(%externref %p2)
+  ret void
+}
diff --git a/llvm/test/Transforms/SLPVectorizer/WebAssembly/externref-no-vectorize.ll b/llvm/test/Transforms/SLPVectorizer/WebAssembly/externref-no-vectorize.ll
new file mode 100644
index 0000000000000..27298fb6ce4e9
--- /dev/null
+++ b/llvm/test/Transforms/SLPVectorizer/WebAssembly/externref-no-vectorize.ll
@@ -0,0 +1,48 @@
+; RUN: opt < %s -passes=slp-vectorizer -mtriple=wasm32-unknown-unknown -S | FileCheck %s
+
+; A WebAssembly reference type cannot be vectorized.
+; This used to crash type legalization while the SLP vectorizer was computing
+; the cost for a <N x externref> value.
+
+target datalayout = "e-m:e-p:32:32-p10:8:8-p20:8:8-i64:64-i128:128-n32:64-S128-ni:1:10:20"
+target triple = "wasm32-unknown-unknown"
+
+%externref = type ptr addrspace(10)
+
+declare %externref @foo()
+declare void @bar(%externref)
+
+; Two externref phis that feed two calls are candidates for being gathered into
+; a vector but they must be left scalar. Derived the following C code:
+;
+;   __externref_t foo(void);
+;   void bar(__externref_t);
+;   void test(int flag, __externref_t ref1, __externref_t ref2) {
+;     if (flag) {
+;       ref1 = foo();
+;       ref2 = foo();
+;     }
+;     bar(ref1); bar(ref2);
+;   }
+;
+; CHECK-LABEL: @test(
+; CHECK-NOT:   phi <{{.*}} x ptr addrspace(10)>
+; CHECK-NOT:   call void @bar(<{{.*}} x ptr addrspace(10)>
+; CHECK:       ret void
+define void @test(i32 %flag, %externref %ref1, %externref %ref2) {
+entry:
+  %c = icmp eq i32 %flag, 0
+  br i1 %c, label %join, label %then
+
+then:
+  %a = tail call %externref @foo()
+  %b = tail call %externref @foo()
+  br label %join
+
+join:
+  %p1 = phi %externref [ %a, %then ], [ %ref1, %entry ]
+  %p2 = phi %externref [ %b, %then ], [ %ref2, %entry ]
+  tail call void @bar(%externref %p1)
+  tail call void @bar(%externref %p2)
+  ret void
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[hoodmane]

I guess maybe instead I should pick up #71540...

[QuantumSegfault]

I guess maybe instead I should pick up #71540...

That probably would help fIx this. Dunno if that would inhibit any optimizations we care about though? Probably not. Just keep it from trying to attempt advanced and not-useful ones like this.

Also note that somebody else already did try to reboot it: #93428

My only concern with doing that right now is that I don't know if TargetExt types can be modeled by LLT. Right now, the use of alternate pointer address spaces is pretty ideal for GlobalISel (which I've been bringing up for WebAssembly). I suppose it's not a huge deal. I don't support externref or funcref anyway in GISel right now anyway, and we can always fall back to SDAG. Use of these things is rare anyway.

And looking at the docs, it seems like moving to TargetExt would bring us closer to being able to represent typed funcref, which is cool.

---

But yeah, I'd like to see that move forward.

[hoodmane]

Okay, I will close this and work on a second reboot of that. Thanks for the guidance @QuantumSegfault.