[WebAssembly] Make RefTypeMem2Local recognize target-features (#88916)

Currently we check `Subtarget->hasReferenceTypes()` to decide whether to
run `RefTypeMem2Local` pass:

https://github.com/llvm/llvm-project/blob/6133878227efc30355c02c2f089e06ce58231a3d/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp#L491-L495

This works fine when `-mattr=+reference-types` is given in the command
line (of `llc` or of `wasm-ld` in case of LTO). This also works fine if
the backend is called by Clang, because Clang's feature set will be
passed to the backend when creating a `TargetMachine`:
https://github.com/llvm/llvm-project/blob/ac791888bbbe58651e597cf7a4b2276424b77a92/clang/lib/CodeGen/BackendUtil.cpp#L549-L550
https://github.com/llvm/llvm-project/blob/ac791888bbbe58651e597cf7a4b2276424b77a92/clang/lib/CodeGen/BackendUtil.cpp#L561-L562

But if the backend compilation is called by `llc`, a `TargetMachine` is
created here:

https://github.com/llvm/llvm-project/blob/bf1ad1d267b1f911cb9846403d2c3d3250a40870/llvm/tools/llc/llc.cpp#L554-L555
And if the backend is called by `wasm-ld`'s LTO, a `TargetMachine` is
created here:

https://github.com/llvm/llvm-project/blob/ac791888bbbe58651e597cf7a4b2276424b77a92/llvm/lib/LTO/LTOBackend.cpp#L513
At this point, in the both places, the created `TargetMachine` only has
access to target features given by the command line with `-mattr=` and
doesn't have access to bitcode functions' `target-features` attribute.

We later gather the target features used by functions and store that
info in the `TargetMachine` in `CoalesceFeaturesAndStripAtomics`,
https://github.com/llvm/llvm-project/blob/ac791888bbbe58651e597cf7a4b2276424b77a92/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp#L202-L206
but this runs in the pass pipeline driven by the pass manager, so this
has not run by the time we check `Subtarget->hasReferenceTypes()` in
`WebAssemblyPassConfig::addISelPrepare`. So currently `RefTypeMem2Local`
would not run on those functions with
`"target-features"="+reference-types"` attributes if the backend is
called by `llc` or `wasm-ld`.

So this makes `RefTypeMem2Local` pass run unconditionally, and checks
`target-featurs` function attribute to decide whether to run the pass on
each function. This allows the pass to run with `wasm-ld` + LTO and
`llc`, even if `-mattr=+reference-types` is not explicitly given in the
command line again, as long as `+reference-types` is in the function's
`target-features` attribute.

This also covers the case we give the target features by the command
line like `llc -mattr=+reference-types` and not in the bitcode
function's attribute, because attributes given in the command line will
be stored in the function's attributes anyway:

https://github.com/llvm/llvm-project/blob/bd28889732e14ac6baca686c3ec99a82fc9cd89d/llvm/lib/CodeGen/CommandFlags.cpp#L673-L674
https://github.com/llvm/llvm-project/blob/bd28889732e14ac6baca686c3ec99a82fc9cd89d/llvm/lib/CodeGen/CommandFlags.cpp#L732-L733

With this PR,
- `lto0.test_externref_emjs`
- `thinlto0.test_externref_emjs`,
- `lto0.test_externref_emjs_dynlink`,
- `thinlto0.test_externref_emjs_dynlnk`

pass. These currently fail but don't get checked in the CI. I think they
used to pass but started to fail after #83196, because we used to run
mem2reg even with `-O0` before that.
(`ltoN` (N > 0) tests are not affected because they run mem2reg anyway
so they don't need `RefTypeMem2Local`)

llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp

@@ -86,6 +86,9 @@ bool WebAssemblyRefTypeMem2Local::runOnFunction(Function &F) {
                        "********** Function: "
                     << F.getName() << '\n');
 
-  visit(F);
+  if (F.getFnAttribute("target-features")
+          .getValueAsString()
+          .contains("+reference-types"))
+    visit(F);
   return Changed;
 }

llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

@@ -484,16 +484,9 @@ void WebAssemblyPassConfig::addIRPasses() {
 }
 
 void WebAssemblyPassConfig::addISelPrepare() {
-  WebAssemblyTargetMachine *WasmTM =
-      static_cast<WebAssemblyTargetMachine *>(TM);
-  const WebAssemblySubtarget *Subtarget =
-      WasmTM->getSubtargetImpl(std::string(WasmTM->getTargetCPU()),
-                               std::string(WasmTM->getTargetFeatureString()));
-  if (Subtarget->hasReferenceTypes()) {
-    // We need to move reference type allocas to WASM_ADDRESS_SPACE_VAR so that
-    // loads and stores are promoted to local.gets/local.sets.
-    addPass(createWebAssemblyRefTypeMem2Local());
-  }
+  // We need to move reference type allocas to WASM_ADDRESS_SPACE_VAR so that
+  // loads and stores are promoted to local.gets/local.sets.
+  addPass(createWebAssemblyRefTypeMem2Local());
   // Lower atomics and TLS if necessary
   addPass(new CoalesceFeaturesAndStripAtomics(&getWebAssemblyTargetMachine()));
 

llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll

@@ -1,4 +1,5 @@
 ; RUN: llc < %s -mattr=+reference-types -stop-after=wasm-ref-type-mem2local | FileCheck %s
+; RUN: llc < %s -stop-after=wasm-ref-type-mem2local | FileCheck %s --check-prefix=ATTR
 
 target triple = "wasm32-unknown-unknown"
 
@@ -51,7 +52,40 @@ entry:
   %i32.loaded = load i32, ptr %alloc.i32
   call void @take_i32(i32 %i32.loaded)
   ; CHECK: %alloc.i32 = alloca i32, align 4{{$}}
-  ; CHECK-NOT: addrspace(1)
+  ; CHECK-NOT: alloca i32 {{.*}} addrspace(1)
 
   ret void
 }
+
+; The same function as test_ref_type_mem2local, but here +reference-types is
+; given in the function attribute.
+; Reference type allocas should be moved to addrspace(1)
+; ATTR-LABEL: @test_ref_type_mem2local_func_attr
+define void @test_ref_type_mem2local_func_attr() #0 {
+entry:
+  %alloc.externref = alloca %externref, align 1
+  %eref = call %externref @get_externref()
+  store %externref %eref, ptr %alloc.externref, align 1
+  %eref.loaded = load %externref, ptr %alloc.externref, align 1
+  call void @take_externref(%externref %eref.loaded)
+  ; ATTR:      %alloc.externref.var = alloca ptr addrspace(10), align 1, addrspace(1)
+  ; ATTR-NEXT: %eref = call ptr addrspace(10) @get_externref()
+  ; ATTR-NEXT: store ptr addrspace(10) %eref, ptr addrspace(1) %alloc.externref.var, align 1
+  ; ATTR-NEXT: %eref.loaded = load ptr addrspace(10), ptr addrspace(1) %alloc.externref.var, align 1
+  ; ATTR-NEXT: call void @take_externref(ptr addrspace(10) %eref.loaded)
+
+  %alloc.funcref = alloca %funcref, align 1
+  %fref = call %funcref @get_funcref()
+  store %funcref %fref, ptr %alloc.funcref, align 1
+  %fref.loaded = load %funcref, ptr %alloc.funcref, align 1
+  call void @take_funcref(%funcref %fref.loaded)
+  ; ATTR-NEXT: %alloc.funcref.var = alloca ptr addrspace(20), align 1, addrspace(1)
+  ; ATTR-NEXT: %fref = call ptr addrspace(20) @get_funcref()
+  ; ATTR-NEXT: store ptr addrspace(20) %fref, ptr addrspace(1) %alloc.funcref.var, align 1
+  ; ATTR-NEXT: %fref.loaded = load ptr addrspace(20), ptr addrspace(1) %alloc.funcref.var, align 1
+  ; ATTR-NEXT: call void @take_funcref(ptr addrspace(20) %fref.loaded)
+
+  ret void
+}
+
+attributes #0 = { "target-features"="+reference-types" }