Skip to content

Add pass to instrument memory.grow instructions #7388

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add pass to instrument memory.grow instructions #7388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 61 additions & 4 deletions src/passes/InstrumentMemory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@

#include "asmjs/shared-constants.h"
#include "shared-constants.h"
#include "support/string.h"
#include <pass.h>
#include <wasm-builder.h>
#include <wasm.h>
Expand Down Expand Up @@ -93,14 +94,26 @@ static Name array_set_val_f32("array_set_val_f32");
static Name array_set_val_f64("array_set_val_f64");
static Name array_get_index("array_get_index");
static Name array_set_index("array_set_index");
static Name memory_grow_pre("memory_grow_pre");
static Name memory_grow_post("memory_grow_post");

// TODO: Add support for atomicRMW/cmpxchg

struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
// Adds calls to new imports.
bool addsEffects() override { return true; }
using InstructionFilter = std::optional<std::unordered_set<std::string>>;

#define CHECK_EXPRESSION(expr) \
do { \
if (filter && !filter->count(expr)) { \
return; \
} \
} while (false)

struct AddInstrumentation : public WalkerPass<PostWalker<AddInstrumentation>> {
explicit AddInstrumentation(InstructionFilter filter)
: filter(std::move(filter)) {}
void visitLoad(Load* curr) {
CHECK_EXPRESSION("load");

id++;
Builder builder(*getModule());
auto mem = getModule()->getMemory(curr->memory);
Expand Down Expand Up @@ -134,6 +147,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

void visitStore(Store* curr) {
CHECK_EXPRESSION("store");

id++;
Builder builder(*getModule());
auto mem = getModule()->getMemory(curr->memory);
Expand Down Expand Up @@ -167,6 +182,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

void visitStructGet(StructGet* curr) {
CHECK_EXPRESSION("struct.get");

Builder builder(*getModule());
Name target;
if (curr->type == Type::i32) {
Expand All @@ -185,6 +202,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

void visitStructSet(StructSet* curr) {
CHECK_EXPRESSION("struct.set");

Builder builder(*getModule());
Name target;
if (curr->value->type == Type::i32) {
Expand All @@ -205,6 +224,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

void visitArrayGet(ArrayGet* curr) {
CHECK_EXPRESSION("array.get");

Builder builder(*getModule());
curr->index =
builder.makeCall(array_get_index,
Expand All @@ -227,6 +248,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

void visitArraySet(ArraySet* curr) {
CHECK_EXPRESSION("array.set");

Builder builder(*getModule());
curr->index =
builder.makeCall(array_set_index,
Expand All @@ -250,10 +273,28 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
curr->value->type);
}

void visitMemoryGrow(MemoryGrow* curr) {
CHECK_EXPRESSION("memory.grow");

id++;
Builder builder(*getModule());
auto addressType = getModule()->getMemory(curr->memory)->addressType;
curr->delta =
builder.makeCall(memory_grow_pre,
{builder.makeConst(int32_t(id)), curr->delta},
addressType);
replaceCurrent(builder.makeCall(
memory_grow_post, {builder.makeConst(int32_t(id)), curr}, addressType));
}

void visitModule(Module* curr) {
auto addressType =
curr->memories.empty() ? Type::i32 : curr->memories[0]->addressType;

// Grow.
addImport(curr, memory_grow_pre, {Type::i32, addressType}, addressType);
addImport(curr, memory_grow_post, {Type::i32, addressType}, addressType);

// Load.
addImport(curr,
load_ptr,
Expand Down Expand Up @@ -300,7 +341,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}

private:
Index id;
Index id = 0;
InstructionFilter filter;

void addImport(Module* curr, Name name, Type params, Type results) {
auto import = Builder::makeFunction(name, Signature(params, results), {});
Expand All @@ -310,6 +352,21 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
}
};

struct InstrumentMemory : Pass {
// Adds calls to new imports.
bool addsEffects() override { return true; }

void run(Module* module) override {
auto arg = getArgumentOrDefault("instrument-memory", "");
InstructionFilter instructions;
if (arg.size() > 0) {
String::Split s(arg, ",");
instructions = std::unordered_set<std::string>{s.begin(), s.end()};
}
AddInstrumentation(std::move(instructions)).run(getPassRunner(), module);
}
};

Pass* createInstrumentMemoryPass() { return new InstrumentMemory(); }

} // namespace wasm
154 changes: 154 additions & 0 deletions test/lit/passes/instrument-memory-filter.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,154 @@
;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.

;; RUN: foreach %s %t wasm-opt --instrument-memory="load,memory.grow" -S -o - | filecheck %s

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a comment explaining that only i32.load and memory.grow should be instrumented below.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see the comment - where is it?

The first function should have a comment "we instrument grow and loads because we were asked to by the flags. The second, "we do not instrument store since we were not asked to by the flags".

;; The test validates the instruction filter used in this pass.
(module
(memory 256 256)
;; CHECK: (type $0 (func (param i32 i32) (result i32)))

;; CHECK: (type $1 (func))
(type $1 (func))
;; CHECK: (type $2 (func (param i32 i32 i32 i32) (result i32)))

;; CHECK: (type $3 (func (param i32 i64) (result i64)))

;; CHECK: (type $4 (func (param i32 f32) (result f32)))

;; CHECK: (type $5 (func (param i32 f64) (result f64)))

;; CHECK: (import "env" "memory_grow_pre" (func $memory_grow_pre (param i32 i32) (result i32)))

;; CHECK: (import "env" "memory_grow_post" (func $memory_grow_post (param i32 i32) (result i32)))

;; CHECK: (import "env" "load_ptr" (func $load_ptr (param i32 i32 i32 i32) (result i32)))

;; CHECK: (import "env" "load_val_i32" (func $load_val_i32 (param i32 i32) (result i32)))

;; CHECK: (import "env" "load_val_i64" (func $load_val_i64 (param i32 i64) (result i64)))

;; CHECK: (import "env" "load_val_f32" (func $load_val_f32 (param i32 f32) (result f32)))

;; CHECK: (import "env" "load_val_f64" (func $load_val_f64 (param i32 f64) (result f64)))

;; CHECK: (import "env" "store_ptr" (func $store_ptr (param i32 i32 i32 i32) (result i32)))

;; CHECK: (import "env" "store_val_i32" (func $store_val_i32 (param i32 i32) (result i32)))

;; CHECK: (import "env" "store_val_i64" (func $store_val_i64 (param i32 i64) (result i64)))

;; CHECK: (import "env" "store_val_f32" (func $store_val_f32 (param i32 f32) (result f32)))

;; CHECK: (import "env" "store_val_f64" (func $store_val_f64 (param i32 f64) (result f64)))

;; CHECK: (memory $0 256 256)

;; CHECK: (func $A
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (call $load_val_i32
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: (i32.load8_s
;; CHECK-NEXT: (call $load_ptr
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (call $load_val_i32
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: (i32.load8_u
;; CHECK-NEXT: (call $load_ptr
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (call $load_val_i64
;; CHECK-NEXT: (i32.const 3)
;; CHECK-NEXT: (i64.load16_s offset=8 align=1
;; CHECK-NEXT: (call $load_ptr
;; CHECK-NEXT: (i32.const 3)
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: (i32.const 8)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (call $load_val_i64
;; CHECK-NEXT: (i32.const 4)
;; CHECK-NEXT: (i64.load32_s offset=10 align=2
;; CHECK-NEXT: (call $load_ptr
;; CHECK-NEXT: (i32.const 4)
;; CHECK-NEXT: (i32.const 4)
;; CHECK-NEXT: (i32.const 10)
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $A (type $1)
;; "*.load*" instructions in this function are instrumented because of the
;; "load" filter included in the command.
(drop (i32.load8_s (i32.const 0)))
(drop (i32.load8_u (i32.const 0)))
(drop (i64.load16_s offset=8 align=1 (i32.const 0)))
(drop (i64.load32_s offset=10 align=2 (i32.const 0)))
)

;; CHECK: (func $B
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (call $memory_grow_post
;; CHECK-NEXT: (i32.const 5)
;; CHECK-NEXT: (memory.grow
;; CHECK-NEXT: (call $memory_grow_pre
;; CHECK-NEXT: (i32.const 5)
;; CHECK-NEXT: (i32.const 4)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $B (type $1)
;; "memory.grow" instructions in this function are instrumented because of the
;; "memory.grow" filter included in the command.
(drop (memory.grow (i32.const 4)))
)

;; CHECK: (func $C
;; CHECK-NEXT: (i32.store8
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (i32.const 1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (i32.store16
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (i32.const 2)
;; CHECK-NEXT: )
;; CHECK-NEXT: (i64.store16 offset=5
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (i64.const 5)
;; CHECK-NEXT: )
;; CHECK-NEXT: (f64.store offset=9 align=2
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: (f64.const 9)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $C (type $1)
;; "*.store*" instructions in this function are not instrumented because the
;; filter is non-empty and doesn't specify "store" instructions.
(i32.store8 (i32.const 0) (i32.const 1))
(i32.store16 (i32.const 0) (i32.const 2))
(i64.store16 offset=5 align=2 (i32.const 0) (i64.const 5))
(f64.store offset=9 align=2 (i32.const 0) (f64.const 9))
)
)
4 changes: 4 additions & 0 deletions test/lit/passes/instrument-memory-gc.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@

;; CHECK: (type $8 (func (param (ref $array))))

;; CHECK: (import "env" "memory_grow_pre" (func $memory_grow_pre (type $0) (param i32 i32) (result i32)))

;; CHECK: (import "env" "memory_grow_post" (func $memory_grow_post (type $0) (param i32 i32) (result i32)))

;; CHECK: (import "env" "load_ptr" (func $load_ptr (type $6) (param i32 i32 i32 i32) (result i32)))

;; CHECK: (import "env" "load_val_i32" (func $load_val_i32 (type $0) (param i32 i32) (result i32)))
Expand Down
Loading
Loading