From d4cdb516eee49ecaf36380af4d8f923cc475e1d7 Mon Sep 17 00:00:00 2001 From: Heejin Ahn Date: Tue, 27 Feb 2024 14:00:43 -0800 Subject: [PATCH] [WebAssembly] Add RefTypeMem2Local pass (#81965) 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 | 57 ++++++++++++ 5 files changed, 152 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..d3c60ee289dfd2 --- /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 { + StringRef getPassName() const override { + return "WebAssembly Reference 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 b2f7ee970a732f..d088c7d925ddf0 100644 --- a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp +++ b/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp @@ -77,6 +77,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..a38243ca218cc1 --- /dev/null +++ b/llvm/test/CodeGen/WebAssembly/ref-type-mem2local.ll @@ -0,0 +1,57 @@ +; 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 %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() { +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() { +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 +}