[WebAssembly] Make RefTypeMem2Local recognize target-features

[aheejin]

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

llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

Lines 491 to 495 in 6133878

	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());
	}

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:

llvm-project/clang/lib/CodeGen/BackendUtil.cpp

Lines 549 to 550 in ac79188

	std::string FeaturesStr =
	llvm::join(TargetOpts.Features.begin(), TargetOpts.Features.end(), ",");

llvm-project/clang/lib/CodeGen/BackendUtil.cpp

Lines 561 to 562 in ac79188

	TM.reset(TheTarget->createTargetMachine(Triple, TargetOpts.CPU, FeaturesStr,
	Options, RM, CM, OptLevel));

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

llvm-project/llvm/tools/llc/llc.cpp

Lines 554 to 555 in bf1ad1d

	Target = std::unique_ptr<TargetMachine>(TheTarget->createTargetMachine(
	TheTriple.getTriple(), CPUStr, FeaturesStr, Options, RM, CM, OLvl));

And if the backend is called by wasm-ld's LTO, a TargetMachine is created here:

llvm-project/llvm/lib/LTO/LTOBackend.cpp

Line 513 in ac79188

std::unique_ptr<TargetMachine> TM = createTargetMachine(C, *TOrErr, Mod);

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,

llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

Lines 202 to 206 in ac79188

	FeatureBitset Features = coalesceFeatures(M);
	std::string FeatureStr =
	getFeatureString(Features, WasmTM->getTargetFeatureString());
	WasmTM->setTargetFeatureString(FeatureStr);

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:

llvm-project/llvm/lib/CodeGen/CommandFlags.cpp

Lines 673 to 674 in bd28889

	if (OldFeatures.empty())
	NewAttrs.addAttribute("target-features", Features);

llvm-project/llvm/lib/CodeGen/CommandFlags.cpp

Lines 732 to 733 in bd28889

	// Let NewAttrs override Attrs.
	F.setAttributes(Attrs.addFnAttributes(Ctx, NewAttrs));

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)

[llvmbot]

@llvm/pr-subscribers-backend-webassembly

Author: Heejin Ahn (aheejin)

Changes

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

llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

Lines 491 to 495 in 6133878

	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());
	}

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:

llvm-project/clang/lib/CodeGen/BackendUtil.cpp

Lines 549 to 550 in ac79188

	std::string FeaturesStr =
	llvm::join(TargetOpts.Features.begin(), TargetOpts.Features.end(), ",");

llvm-project/clang/lib/CodeGen/BackendUtil.cpp

Lines 561 to 562 in ac79188

	TM.reset(TheTarget->createTargetMachine(Triple, TargetOpts.CPU, FeaturesStr,
	Options, RM, CM, OptLevel));

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

llvm-project/llvm/tools/llc/llc.cpp

Lines 554 to 555 in bf1ad1d

	Target = std::unique_ptr<TargetMachine>(TheTarget->createTargetMachine(
	TheTriple.getTriple(), CPUStr, FeaturesStr, Options, RM, CM, OLvl));

And if the backend is called by wasm-ld's LTO, a TargetMachine is created here:

llvm-project/llvm/lib/LTO/LTOBackend.cpp

Line 513 in ac79188

std::unique_ptr<TargetMachine> TM = createTargetMachine(C, *TOrErr, Mod);

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,

llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

Lines 202 to 206 in ac79188

	FeatureBitset Features = coalesceFeatures(M);
	std::string FeatureStr =
	getFeatureString(Features, WasmTM->getTargetFeatureString());
	WasmTM->setTargetFeatureString(FeatureStr);

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:

llvm-project/llvm/lib/CodeGen/CommandFlags.cpp

Lines 673 to 674 in bd28889

	if (OldFeatures.empty())
	NewAttrs.addAttribute("target-features", Features);

llvm-project/llvm/lib/CodeGen/CommandFlags.cpp

Lines 732 to 733 in bd28889

	// Let NewAttrs override Attrs.
	F.setAttributes(Attrs.addFnAttributes(Ctx, NewAttrs));

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)

---

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

3 Files Affected:

• (modified) llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp (+4-1)
• (modified) llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp (+3-10)
• (added) llvm/test/CodeGen/WebAssembly/ref-type-mem2local-func-attr.ll (+63)

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
index d3c60ee289dfd2..04b4c7d78aabb3 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
+++ b/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;
 }
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
index 769ee765e19078..61da0ecd0f5d1c 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
+++ b/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()));
 
diff --git a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local-func-attr.ll b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local-func-attr.ll
new file mode 100644
index 00000000000000..f547d1393e26be
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local-func-attr.ll
@@ -0,0 +1,63 @@
+; RUN: llc < %s -stop-after=wasm-ref-type-mem2local | FileCheck %s
+
+; This file is the same as ref-type-mem2local.ll, except we pass
+; '+reference-types' target feature using a function attribute instead of
+; '-mattr=+reference-types' command line argument.
+
+target triple = "wasm32-unknown-unknown"
+
+%externref = type ptr addrspace(10)
+%funcref = type ptr addrspace(20)
+
+declare %externref @get_externref()
+declare %funcref @get_funcref()
+declare i32 @get_i32()
+declare void @take_externref(%externref)
+declare void @take_funcref(%funcref)
+declare void @take_i32(i32)
+
+; Reference type allocas should be moved to addrspace(1)
+; CHECK-LABEL: @test_ref_type_mem2local
+define void @test_ref_type_mem2local() #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)
+  ; CHECK:      %alloc.externref.var = alloca ptr addrspace(10), align 1, addrspace(1)
+  ; CHECK-NEXT: %eref = call ptr addrspace(10) @get_externref()
+  ; CHECK-NEXT: store ptr addrspace(10) %eref, ptr addrspace(1) %alloc.externref.var, align 1
+  ; CHECK-NEXT: %eref.loaded = load ptr addrspace(10), ptr addrspace(1) %alloc.externref.var, align 1
+  ; CHECK-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)
+  ; CHECK-NEXT: %alloc.funcref.var = alloca ptr addrspace(20), align 1, addrspace(1)
+  ; CHECK-NEXT: %fref = call ptr addrspace(20) @get_funcref()
+  ; CHECK-NEXT: store ptr addrspace(20) %fref, ptr addrspace(1) %alloc.funcref.var, align 1
+  ; CHECK-NEXT: %fref.loaded = load ptr addrspace(20), ptr addrspace(1) %alloc.funcref.var, align 1
+  ; CHECK-NEXT: call void @take_funcref(ptr addrspace(20) %fref.loaded)
+
+  ret void
+}
+
+; POD type allocas should stay the same
+; CHECK-LABEL: @test_pod_type
+define void @test_pod_type() #0 {
+entry:
+  %alloc.i32 = alloca i32
+  %i32 = call i32 @get_i32()
+  store i32 %i32, ptr %alloc.i32
+  %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)
+
+  ret void
+}
+
+attributes #0 = { "target-features"="+reference-types" }