Skip to content

Commit cc1cc3d

Browse files
loganekloganek
authored
InstrumentMemory: Allow filtering by instruction, and instrument memory.grow (#7388)
1 parent ada1581 commit cc1cc3d

File tree

5 files changed

+268
-11
lines changed

5 files changed

+268
-11
lines changed

‎src/passes/InstrumentMemory.cpp

Lines changed: 61 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@
5959

6060
#include "asmjs/shared-constants.h"
6161
#include "shared-constants.h"
62+
#include "support/string.h"
6263
#include <pass.h>
6364
#include <wasm-builder.h>
6465
#include <wasm.h>
@@ -93,14 +94,26 @@ static Name array_set_val_f32("array_set_val_f32");
9394
static Name array_set_val_f64("array_set_val_f64");
9495
static Name array_get_index("array_get_index");
9596
static Name array_set_index("array_set_index");
97+
static Name memory_grow_pre("memory_grow_pre");
98+
static Name memory_grow_post("memory_grow_post");
9699

97100
// TODO: Add support for atomicRMW/cmpxchg
98101

99-
struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
100-
// Adds calls to new imports.
101-
bool addsEffects() override { return true; }
102+
using InstructionFilter = std::optional<std::unordered_set<std::string>>;
103+
104+
#define CHECK_EXPRESSION(expr) \
105+
do { \
106+
if (filter && !filter->count(expr)) { \
107+
return; \
108+
} \
109+
} while (false)
102110

111+
struct AddInstrumentation : public WalkerPass<PostWalker<AddInstrumentation>> {
112+
explicit AddInstrumentation(InstructionFilter filter)
113+
: filter(std::move(filter)) {}
103114
void visitLoad(Load* curr) {
115+
CHECK_EXPRESSION("load");
116+
104117
id++;
105118
Builder builder(*getModule());
106119
auto mem = getModule()->getMemory(curr->memory);
@@ -134,6 +147,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
134147
}
135148

136149
void visitStore(Store* curr) {
150+
CHECK_EXPRESSION("store");
151+
137152
id++;
138153
Builder builder(*getModule());
139154
auto mem = getModule()->getMemory(curr->memory);
@@ -167,6 +182,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
167182
}
168183

169184
void visitStructGet(StructGet* curr) {
185+
CHECK_EXPRESSION("struct.get");
186+
170187
Builder builder(*getModule());
171188
Name target;
172189
if (curr->type == Type::i32) {
@@ -185,6 +202,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
185202
}
186203

187204
void visitStructSet(StructSet* curr) {
205+
CHECK_EXPRESSION("struct.set");
206+
188207
Builder builder(*getModule());
189208
Name target;
190209
if (curr->value->type == Type::i32) {
@@ -205,6 +224,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
205224
}
206225

207226
void visitArrayGet(ArrayGet* curr) {
227+
CHECK_EXPRESSION("array.get");
228+
208229
Builder builder(*getModule());
209230
curr->index =
210231
builder.makeCall(array_get_index,
@@ -227,6 +248,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
227248
}
228249

229250
void visitArraySet(ArraySet* curr) {
251+
CHECK_EXPRESSION("array.set");
252+
230253
Builder builder(*getModule());
231254
curr->index =
232255
builder.makeCall(array_set_index,
@@ -250,10 +273,28 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
250273
curr->value->type);
251274
}
252275

276+
void visitMemoryGrow(MemoryGrow* curr) {
277+
CHECK_EXPRESSION("memory.grow");
278+
279+
id++;
280+
Builder builder(*getModule());
281+
auto addressType = getModule()->getMemory(curr->memory)->addressType;
282+
curr->delta =
283+
builder.makeCall(memory_grow_pre,
284+
{builder.makeConst(int32_t(id)), curr->delta},
285+
addressType);
286+
replaceCurrent(builder.makeCall(
287+
memory_grow_post, {builder.makeConst(int32_t(id)), curr}, addressType));
288+
}
289+
253290
void visitModule(Module* curr) {
254291
auto addressType =
255292
curr->memories.empty() ? Type::i32 : curr->memories[0]->addressType;
256293

294+
// Grow.
295+
addImport(curr, memory_grow_pre, {Type::i32, addressType}, addressType);
296+
addImport(curr, memory_grow_post, {Type::i32, addressType}, addressType);
297+
257298
// Load.
258299
addImport(curr,
259300
load_ptr,
@@ -300,7 +341,8 @@ struct InstrumentMemory : public WalkerPass<PostWalker<InstrumentMemory>> {
300341
}
301342

302343
private:
303-
Index id;
344+
Index id = 0;
345+
InstructionFilter filter;
304346

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

355+
struct InstrumentMemory : Pass {
356+
// Adds calls to new imports.
357+
bool addsEffects() override { return true; }
358+
359+
void run(Module* module) override {
360+
auto arg = getArgumentOrDefault("instrument-memory", "");
361+
InstructionFilter instructions;
362+
if (arg.size() > 0) {
363+
String::Split s(arg, ",");
364+
instructions = std::unordered_set<std::string>{s.begin(), s.end()};
365+
}
366+
AddInstrumentation(std::move(instructions)).run(getPassRunner(), module);
367+
}
368+
};
369+
313370
Pass* createInstrumentMemoryPass() { return new InstrumentMemory(); }
314371

315372
} // namespace wasm

‎test/lit/passes/instrument-memory-filter.wast

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
2+
3+
;; RUN: foreach %s %t wasm-opt --instrument-memory="load,memory.grow" -S -o - | filecheck %s
4+
5+
;; The test validates the instruction filter used in this pass.
6+
(module
7+
(memory 256 256)
8+
;; CHECK: (type $0 (func (param i32 i32) (result i32)))
9+
10+
;; CHECK: (type $1 (func))
11+
(type $1 (func))
12+
;; CHECK: (type $2 (func (param i32 i32 i32 i32) (result i32)))
13+
14+
;; CHECK: (type $3 (func (param i32 i64) (result i64)))
15+
16+
;; CHECK: (type $4 (func (param i32 f32) (result f32)))
17+
18+
;; CHECK: (type $5 (func (param i32 f64) (result f64)))
19+
20+
;; CHECK: (import "env" "memory_grow_pre" (func $memory_grow_pre (param i32 i32) (result i32)))
21+
22+
;; CHECK: (import "env" "memory_grow_post" (func $memory_grow_post (param i32 i32) (result i32)))
23+
24+
;; CHECK: (import "env" "load_ptr" (func $load_ptr (param i32 i32 i32 i32) (result i32)))
25+
26+
;; CHECK: (import "env" "load_val_i32" (func $load_val_i32 (param i32 i32) (result i32)))
27+
28+
;; CHECK: (import "env" "load_val_i64" (func $load_val_i64 (param i32 i64) (result i64)))
29+
30+
;; CHECK: (import "env" "load_val_f32" (func $load_val_f32 (param i32 f32) (result f32)))
31+
32+
;; CHECK: (import "env" "load_val_f64" (func $load_val_f64 (param i32 f64) (result f64)))
33+
34+
;; CHECK: (import "env" "store_ptr" (func $store_ptr (param i32 i32 i32 i32) (result i32)))
35+
36+
;; CHECK: (import "env" "store_val_i32" (func $store_val_i32 (param i32 i32) (result i32)))
37+
38+
;; CHECK: (import "env" "store_val_i64" (func $store_val_i64 (param i32 i64) (result i64)))
39+
40+
;; CHECK: (import "env" "store_val_f32" (func $store_val_f32 (param i32 f32) (result f32)))
41+
42+
;; CHECK: (import "env" "store_val_f64" (func $store_val_f64 (param i32 f64) (result f64)))
43+
44+
;; CHECK: (memory $0 256 256)
45+
46+
;; CHECK: (func $A
47+
;; CHECK-NEXT: (drop
48+
;; CHECK-NEXT: (call $load_val_i32
49+
;; CHECK-NEXT: (i32.const 1)
50+
;; CHECK-NEXT: (i32.load8_s
51+
;; CHECK-NEXT: (call $load_ptr
52+
;; CHECK-NEXT: (i32.const 1)
53+
;; CHECK-NEXT: (i32.const 1)
54+
;; CHECK-NEXT: (i32.const 0)
55+
;; CHECK-NEXT: (i32.const 0)
56+
;; CHECK-NEXT: )
57+
;; CHECK-NEXT: )
58+
;; CHECK-NEXT: )
59+
;; CHECK-NEXT: )
60+
;; CHECK-NEXT: (drop
61+
;; CHECK-NEXT: (call $load_val_i32
62+
;; CHECK-NEXT: (i32.const 2)
63+
;; CHECK-NEXT: (i32.load8_u
64+
;; CHECK-NEXT: (call $load_ptr
65+
;; CHECK-NEXT: (i32.const 2)
66+
;; CHECK-NEXT: (i32.const 1)
67+
;; CHECK-NEXT: (i32.const 0)
68+
;; CHECK-NEXT: (i32.const 0)
69+
;; CHECK-NEXT: )
70+
;; CHECK-NEXT: )
71+
;; CHECK-NEXT: )
72+
;; CHECK-NEXT: )
73+
;; CHECK-NEXT: (drop
74+
;; CHECK-NEXT: (call $load_val_i64
75+
;; CHECK-NEXT: (i32.const 3)
76+
;; CHECK-NEXT: (i64.load16_s offset=8 align=1
77+
;; CHECK-NEXT: (call $load_ptr
78+
;; CHECK-NEXT: (i32.const 3)
79+
;; CHECK-NEXT: (i32.const 2)
80+
;; CHECK-NEXT: (i32.const 8)
81+
;; CHECK-NEXT: (i32.const 0)
82+
;; CHECK-NEXT: )
83+
;; CHECK-NEXT: )
84+
;; CHECK-NEXT: )
85+
;; CHECK-NEXT: )
86+
;; CHECK-NEXT: (drop
87+
;; CHECK-NEXT: (call $load_val_i64
88+
;; CHECK-NEXT: (i32.const 4)
89+
;; CHECK-NEXT: (i64.load32_s offset=10 align=2
90+
;; CHECK-NEXT: (call $load_ptr
91+
;; CHECK-NEXT: (i32.const 4)
92+
;; CHECK-NEXT: (i32.const 4)
93+
;; CHECK-NEXT: (i32.const 10)
94+
;; CHECK-NEXT: (i32.const 0)
95+
;; CHECK-NEXT: )
96+
;; CHECK-NEXT: )
97+
;; CHECK-NEXT: )
98+
;; CHECK-NEXT: )
99+
;; CHECK-NEXT: )
100+
(func $A (type $1)
101+
;; "*.load*" instructions in this function are instrumented because of the
102+
;; "load" filter included in the command.
103+
(drop (i32.load8_s (i32.const 0)))
104+
(drop (i32.load8_u (i32.const 0)))
105+
(drop (i64.load16_s offset=8 align=1 (i32.const 0)))
106+
(drop (i64.load32_s offset=10 align=2 (i32.const 0)))
107+
)
108+
109+
;; CHECK: (func $B
110+
;; CHECK-NEXT: (drop
111+
;; CHECK-NEXT: (call $memory_grow_post
112+
;; CHECK-NEXT: (i32.const 5)
113+
;; CHECK-NEXT: (memory.grow
114+
;; CHECK-NEXT: (call $memory_grow_pre
115+
;; CHECK-NEXT: (i32.const 5)
116+
;; CHECK-NEXT: (i32.const 4)
117+
;; CHECK-NEXT: )
118+
;; CHECK-NEXT: )
119+
;; CHECK-NEXT: )
120+
;; CHECK-NEXT: )
121+
;; CHECK-NEXT: )
122+
(func $B (type $1)
123+
;; "memory.grow" instructions in this function are instrumented because of the
124+
;; "memory.grow" filter included in the command.
125+
(drop (memory.grow (i32.const 4)))
126+
)
127+
128+
;; CHECK: (func $C
129+
;; CHECK-NEXT: (i32.store8
130+
;; CHECK-NEXT: (i32.const 0)
131+
;; CHECK-NEXT: (i32.const 1)
132+
;; CHECK-NEXT: )
133+
;; CHECK-NEXT: (i32.store16
134+
;; CHECK-NEXT: (i32.const 0)
135+
;; CHECK-NEXT: (i32.const 2)
136+
;; CHECK-NEXT: )
137+
;; CHECK-NEXT: (i64.store16 offset=5
138+
;; CHECK-NEXT: (i32.const 0)
139+
;; CHECK-NEXT: (i64.const 5)
140+
;; CHECK-NEXT: )
141+
;; CHECK-NEXT: (f64.store offset=9 align=2
142+
;; CHECK-NEXT: (i32.const 0)
143+
;; CHECK-NEXT: (f64.const 9)
144+
;; CHECK-NEXT: )
145+
;; CHECK-NEXT: )
146+
(func $C (type $1)
147+
;; "*.store*" instructions in this function are not instrumented because the
148+
;; filter is non-empty and doesn't specify "store" instructions.
149+
(i32.store8 (i32.const 0) (i32.const 1))
150+
(i32.store16 (i32.const 0) (i32.const 2))
151+
(i64.store16 offset=5 align=2 (i32.const 0) (i64.const 5))
152+
(f64.store offset=9 align=2 (i32.const 0) (f64.const 9))
153+
)
154+
)

‎test/lit/passes/instrument-memory-gc.wast

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@
2626

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

29+
;; CHECK: (import "env" "memory_grow_pre" (func $memory_grow_pre (type $0) (param i32 i32) (result i32)))
30+
31+
;; CHECK: (import "env" "memory_grow_post" (func $memory_grow_post (type $0) (param i32 i32) (result i32)))
32+
2933
;; CHECK: (import "env" "load_ptr" (func $load_ptr (type $6) (param i32 i32 i32 i32) (result i32)))
3034

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

‎test/lit/passes/instrument-memory.wast

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,21 @@
55

66
(module
77
(memory 256 256)
8+
;; CHECK: (type $0 (func (param i32 i32) (result i32)))
9+
810
;; CHECK: (type $1 (func))
911
(type $1 (func))
1012
;; CHECK: (type $2 (func (param i32 i32 i32 i32) (result i32)))
1113

12-
;; CHECK: (type $3 (func (param i32 i32) (result i32)))
14+
;; CHECK: (type $3 (func (param i32 i64) (result i64)))
15+
16+
;; CHECK: (type $4 (func (param i32 f32) (result f32)))
1317

14-
;; CHECK: (type $4 (func (param i32 i64) (result i64)))
18+
;; CHECK: (type $5 (func (param i32 f64) (result f64)))
1519

16-
;; CHECK: (type $5 (func (param i32 f32) (result f32)))
20+
;; CHECK: (import "env" "memory_grow_pre" (func $memory_grow_pre (param i32 i32) (result i32)))
1721

18-
;; CHECK: (type $6 (func (param i32 f64) (result f64)))
22+
;; CHECK: (import "env" "memory_grow_post" (func $memory_grow_post (param i32 i32) (result i32)))
1923

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

@@ -676,4 +680,21 @@
676680
(f32.store offset=8 align=2 (i32.const 0) (f32.const 8))
677681
(f64.store offset=9 align=2 (i32.const 0) (f64.const 9))
678682
)
683+
684+
;; CHECK: (func $C
685+
;; CHECK-NEXT: (drop
686+
;; CHECK-NEXT: (call $memory_grow_post
687+
;; CHECK-NEXT: (i32.const 47)
688+
;; CHECK-NEXT: (memory.grow
689+
;; CHECK-NEXT: (call $memory_grow_pre
690+
;; CHECK-NEXT: (i32.const 47)
691+
;; CHECK-NEXT: (i32.const 4)
692+
;; CHECK-NEXT: )
693+
;; CHECK-NEXT: )
694+
;; CHECK-NEXT: )
695+
;; CHECK-NEXT: )
696+
;; CHECK-NEXT: )
697+
(func $C (type $1)
698+
(drop (memory.grow (i32.const 4)))
699+
)
679700
)

0 commit comments

Comments
 (0)