[llvm] [WebAssembly] Add RefTypeMem2Local pass (PR #81965)

Heejin Ahn via llvm-commits llvm-commits at lists.llvm.org
Thu Feb 22 16:39:08 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/5] [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/5] 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/5] 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
+}

>From b1ec1f76dff7ca8feb2bfbbacbd6ee77b2ed600a Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 16 Feb 2024 18:20:09 +0000
Subject: [PATCH 4/5] change order

---
 llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
index ace39d5573a4e3..f1eda8443b563c 100644
--- a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
+++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
@@ -2,8 +2,8 @@
 
 target triple = "wasm32-unknown-unknown"
 
-%externref = type ptr addrspace(10)
 %funcref = type ptr addrspace(20)
+%externref = type ptr addrspace(10)
 
 declare %funcref @get_funcref()
 declare %externref @get_externref()

>From 9831fb3b4bd20f17c5aa8ef9e7bb2f9c3ca51c9f Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 16 Feb 2024 22:09:16 +0000
Subject: [PATCH 5/5] Cosmetic order change nevermind

---
 llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
index f1eda8443b563c..a38243ca218cc1 100644
--- a/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
+++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll
@@ -2,14 +2,14 @@
 
 target triple = "wasm32-unknown-unknown"
 
-%funcref = type ptr addrspace(20)
 %externref = type ptr addrspace(10)
+%funcref = type ptr addrspace(20)
 
-declare %funcref @get_funcref()
 declare %externref @get_externref()
+declare %funcref @get_funcref()
 declare i32 @get_i32()
-declare void @take_funcref(%funcref)
 declare void @take_externref(%externref)
+declare void @take_funcref(%funcref)
 declare void @take_i32(i32)
 
 ; Reference type allocas should be moved to addrspace(1)



More information about the llvm-commits mailing list