[WebAssembly] Make RefTypeMem2Local recognize target-features

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

@@ -483,16 +483,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" }