[llvm] [WebAssembly] Add RefTypeMem2Local pass (PR #81965)
Heejin Ahn via llvm-commits
llvm-commits at lists.llvm.org
Fri Feb 16 10:04:36 PST 2024
https://github.com/aheejin updated https://github.com/llvm/llvm-project/pull/81965
>From 58beb2afb014e19757b2b468042a3cf7b099f8d5 Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 16 Feb 2024 01:42:38 +0000
Subject: [PATCH 1/3] [WebAssembly] Add RefTypeMem2Local pass
This adds `WebAssemblyRefTypeMem2Local` pass, which changes the address
spaces of reference type `alloca`s to `addrspace(1)`. This in turn
changes the address spaces of all `load` and `store` instructions that
use the `alloca`s.
`addrspace(1)` is `WASM_ADDRESS_SPACE_VAR`, and loads and stores to this
address space become `local.get`s and `local.set`s, thanks to the Wasm
local IR support added in
https://github.com/llvm/llvm-project/commit/82f92e35c6464e23859c29422956caaceb623967.
In a follow-up PR, I am planning to replace the usage of mem2reg pass
with this to solve the reference type `alloca` problems described in
#81575.
---
llvm/lib/Target/WebAssembly/CMakeLists.txt | 1 +
llvm/lib/Target/WebAssembly/WebAssembly.h | 2 +
.../WebAssemblyRefTypeMem2Local.cpp | 91 +++++++++++++++++++
.../WebAssembly/WebAssemblyTargetMachine.cpp | 1 +
.../CodeGen/WebAssembly/ref-type-mem2local.ll | 39 ++++++++
5 files changed, 134 insertions(+)
create mode 100644 llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
create mode 100644 llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
diff --git a/llvm/lib/Target/WebAssembly/CMakeLists.txt b/llvm/lib/Target/WebAssembly/CMakeLists.txt
index bb2ccea5c14598..f430be2653b4ee 100644
--- a/llvm/lib/Target/WebAssembly/CMakeLists.txt
+++ b/llvm/lib/Target/WebAssembly/CMakeLists.txt
@@ -43,6 +43,7 @@ add_llvm_target(WebAssemblyCodeGen
WebAssemblyOptimizeLiveIntervals.cpp
WebAssemblyOptimizeReturned.cpp
WebAssemblyPeephole.cpp
+ WebAssemblyRefTypeMem2Local.cpp
WebAssemblyRegisterInfo.cpp
WebAssemblyRegColoring.cpp
WebAssemblyRegNumbering.cpp
diff --git a/llvm/lib/Target/WebAssembly/WebAssembly.h b/llvm/lib/Target/WebAssembly/WebAssembly.h
index 91765ad117bdb0..1c40addb6d6f78 100644
--- a/llvm/lib/Target/WebAssembly/WebAssembly.h
+++ b/llvm/lib/Target/WebAssembly/WebAssembly.h
@@ -30,6 +30,7 @@ ModulePass *createWebAssemblyAddMissingPrototypes();
ModulePass *createWebAssemblyFixFunctionBitcasts();
FunctionPass *createWebAssemblyOptimizeReturned();
FunctionPass *createWebAssemblyLowerRefTypesIntPtrConv();
+FunctionPass *createWebAssemblyRefTypeMem2Local();
// ISel and immediate followup passes.
FunctionPass *createWebAssemblyISelDag(WebAssemblyTargetMachine &TM,
@@ -59,6 +60,7 @@ ModulePass *createWebAssemblyMCLowerPrePass();
// PassRegistry initialization declarations.
void initializeFixFunctionBitcastsPass(PassRegistry &);
void initializeOptimizeReturnedPass(PassRegistry &);
+void initializeWebAssemblyRefTypeMem2LocalPass(PassRegistry &);
void initializeWebAssemblyAddMissingPrototypesPass(PassRegistry &);
void initializeWebAssemblyArgumentMovePass(PassRegistry &);
void initializeWebAssemblyCFGSortPass(PassRegistry &);
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
new file mode 100644
index 00000000000000..45773fd179184d
--- /dev/null
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
@@ -0,0 +1,91 @@
+//=== WebAssemblyRefTypeMem2Local.cpp - WebAssembly RefType Mem2Local -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Assign reference type allocas to local addrspace (addrspace(1)) so that
+/// their loads and stores can be lowered to local.gets/local.sets.
+///
+//===----------------------------------------------------------------------===//
+
+#include "Utils/WasmAddressSpaces.h"
+#include "Utils/WebAssemblyTypeUtilities.h"
+#include "WebAssembly.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/ValueHandle.h"
+#include "llvm/Pass.h"
+using namespace llvm;
+
+#define DEBUG_TYPE "wasm-ref-type-mem2local"
+
+namespace {
+class WebAssemblyRefTypeMem2Local final
+ : public FunctionPass,
+ public InstVisitor<WebAssemblyRefTypeMem2Local> {
+ StringRef getPassName() const override {
+ return "WebAssembly Refernce Types Memory to Local";
+ }
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.setPreservesCFG();
+ FunctionPass::getAnalysisUsage(AU);
+ }
+
+ bool runOnFunction(Function &F) override;
+ bool Changed = false;
+
+public:
+ static char ID;
+ WebAssemblyRefTypeMem2Local() : FunctionPass(ID) {}
+
+ void visitAllocaInst(AllocaInst &AI);
+};
+} // End anonymous namespace
+
+char WebAssemblyRefTypeMem2Local::ID = 0;
+INITIALIZE_PASS(WebAssemblyRefTypeMem2Local, DEBUG_TYPE,
+ "Assign reference type allocas to local address space", true,
+ false)
+
+FunctionPass *llvm::createWebAssemblyRefTypeMem2Local() {
+ return new WebAssemblyRefTypeMem2Local();
+}
+
+void WebAssemblyRefTypeMem2Local::visitAllocaInst(AllocaInst &AI) {
+ if (WebAssembly::isWebAssemblyReferenceType(AI.getAllocatedType())) {
+ Changed = true;
+ IRBuilder<> IRB(AI.getContext());
+ IRB.SetInsertPoint(&AI);
+ auto *NewAI = IRB.CreateAlloca(AI.getAllocatedType(),
+ WebAssembly::WASM_ADDRESS_SPACE_VAR, nullptr,
+ AI.getName() + ".var");
+
+ // The below is basically equivalent to AI.replaceAllUsesWith(NewAI), but we
+ // cannot use it because it requires the old and new types be the same,
+ // which is not true here because the address spaces are different.
+ if (AI.hasValueHandle())
+ ValueHandleBase::ValueIsRAUWd(&AI, NewAI);
+ if (AI.isUsedByMetadata())
+ ValueAsMetadata::handleRAUW(&AI, NewAI);
+ while (!AI.materialized_use_empty()) {
+ Use &U = *AI.materialized_use_begin();
+ U.set(NewAI);
+ }
+
+ AI.eraseFromParent();
+ }
+}
+
+bool WebAssemblyRefTypeMem2Local::runOnFunction(Function &F) {
+ LLVM_DEBUG(dbgs() << "********** WebAssembly RefType Mem2Local **********\n"
+ "********** Function: "
+ << F.getName() << '\n');
+
+ visit(F);
+ return Changed;
+}
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
index 42043a7b8680a4..681e39e5aa5e32 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
@@ -68,6 +68,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeWebAssemblyTarget() {
initializeLowerGlobalDtorsLegacyPassPass(PR);
initializeFixFunctionBitcastsPass(PR);
initializeOptimizeReturnedPass(PR);
+ initializeWebAssemblyRefTypeMem2LocalPass(PR);
initializeWebAssemblyArgumentMovePass(PR);
initializeWebAssemblySetP2AlignOperandsPass(PR);
initializeWebAssemblyReplacePhysRegsPass(PR);
diff --git a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
new file mode 100644
index 00000000000000..8f09b25cca8d05
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
@@ -0,0 +1,39 @@
+; RUN: opt < %s -wasm-ref-type-mem2local -S | FileCheck %s
+
+target triple = "wasm32-unknown-unknown"
+
+%externref = type ptr addrspace(10)
+%funcref = type ptr addrspace(20)
+
+declare %funcref @get_funcref()
+declare %externref @get_externref()
+declare void @take_funcref(%funcref)
+declare void @take_externref(%externref)
+
+; CHECK-LABEL: @test_ref_type_mem2local
+define void @test_ref_type_mem2local() {
+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
+}
>From d5493d90c5ae786e7fb9d262966923f7d9133330 Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 16 Feb 2024 09:53:03 -0800
Subject: [PATCH 2/3] Update
llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
Co-authored-by: xortoast <xortoast at proton.me>
---
llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
index 45773fd179184d..d3c60ee289dfd2 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyRefTypeMem2Local.cpp
@@ -28,7 +28,7 @@ class WebAssemblyRefTypeMem2Local final
: public FunctionPass,
public InstVisitor<WebAssemblyRefTypeMem2Local> {
StringRef getPassName() const override {
- return "WebAssembly Refernce Types Memory to Local";
+ return "WebAssembly Reference Types Memory to Local";
}
void getAnalysisUsage(AnalysisUsage &AU) const override {
>From a2b29e37d22821c11de5b2be7cc83fb7099bb04f Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 16 Feb 2024 18:04:13 +0000
Subject: [PATCH 3/3] Test i32 allocas as a control group
---
.../CodeGen/WebAssembly/ref-type-mem2local.ll | 18 ++++++++++++++++++
1 file changed, 18 insertions(+)
diff --git a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
index 8f09b25cca8d05..ace39d5573a4e3 100644
--- a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
+++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
@@ -7,9 +7,12 @@ target triple = "wasm32-unknown-unknown"
declare %funcref @get_funcref()
declare %externref @get_externref()
+declare i32 @get_i32()
declare void @take_funcref(%funcref)
declare void @take_externref(%externref)
+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() {
entry:
@@ -37,3 +40,18 @@ entry:
ret void
}
+
+; POD type allocas should stay the same
+; CHECK-LABEL: @test_pod_type
+define void @test_pod_type() {
+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
+}
More information about the llvm-commits
mailing list