Skip to content

Commit

Permalink
Heap2Local: Optimize packed fields (#6480)
Browse files Browse the repository at this point in the history 
Previously we did not optimize a struct or an array with a packed field. As a result a
single packed field in a struct prevented the entire struct from being localized,
which this fixes. This is also useful for arrays as packed arrays are common (e.g. for
string data).
  • Loading branch information
@kripken
kripken authored Apr 9, 2024
1 parent fca1d3a commit 738e8fc
Show file tree
Hide file tree
Showing 3 changed files with 261 additions and 41 deletions.
38 changes: 26 additions & 12 deletions src/passes/Heap2Local.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@
// This optimization focuses on such cases.
//

#include "ir/bits.h"
#include "ir/branch-utils.h"
#include "ir/find_all.h"
#include "ir/local-graph.h"
Expand Down Expand Up @@ -648,6 +649,21 @@ struct Struct2Local : PostWalker<Struct2Local> {
curr->finalize();
}

// Add a mask for packed fields. We add masks on sets rather than on gets
// because gets tend to be more numerous both in code appearances and in
// runtime execution. As a result of masking on sets, the value in the local
// is always the masked value (which is also nice for debugging,
// incidentally).
Expression* addMask(Expression* value, const Field& field) {
if (!field.isPacked()) {
return value;
}

auto mask = Bits::lowBitMask(field.getByteSize() * 8);
return builder.makeBinary(
AndInt32, value, builder.makeConst(int32_t(mask)));
}

void visitStructNew(StructNew* curr) {
if (curr != allocation) {
return;
Expand Down Expand Up @@ -693,9 +709,10 @@ struct Struct2Local : PostWalker<Struct2Local> {

// Copy them to the normal ones.
for (Index i = 0; i < tempIndexes.size(); i++) {
contents.push_back(builder.makeLocalSet(
localIndexes[i],
builder.makeLocalGet(tempIndexes[i], fields[i].type)));
auto* value = builder.makeLocalGet(tempIndexes[i], fields[i].type);
// Add a mask on the values we write.
contents.push_back(
builder.makeLocalSet(localIndexes[i], addMask(value, fields[i])));
}

// TODO Check if the nondefault case does not increase code size in some
Expand All @@ -704,8 +721,11 @@ struct Struct2Local : PostWalker<Struct2Local> {
// defaults.
} else {
// Set the default values.
//
// Note that we must assign the defaults because we might be in a loop,
// that is, there might be a previous value.
//
// Note there is no need to mask as these are zeros anyhow.
for (Index i = 0; i < localIndexes.size(); i++) {
contents.push_back(builder.makeLocalSet(
localIndexes[i],
Expand Down Expand Up @@ -766,7 +786,8 @@ struct Struct2Local : PostWalker<Struct2Local> {
// write the data to the local instead of the heap allocation.
replaceCurrent(builder.makeSequence(
builder.makeDrop(curr->ref),
builder.makeLocalSet(localIndexes[curr->index], curr->value)));
builder.makeLocalSet(localIndexes[curr->index],
addMask(curr->value, fields[curr->index]))));
}

void visitStructGet(StructGet* curr) {
Expand Down Expand Up @@ -1111,14 +1132,7 @@ struct Heap2Local {
}

bool canHandleAsLocal(const Field& field) {
if (!TypeUpdating::canHandleAsLocal(field.type)) {
return false;
}
if (field.isPacked()) {
// TODO: support packed fields by adding coercions/truncations.
return false;
}
return true;
return TypeUpdating::canHandleAsLocal(field.type);
}

bool canHandleAsLocals(Type type) {
Expand Down
236 changes: 226 additions & 10 deletions test/lit/passes/heap2local.wast
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,17 @@

;; CHECK: (type $6 (func (result anyref)))

;; CHECK: (type $struct.packed (struct (field (mut i8))))
(type $struct.packed (struct (field (mut i8))))
;; CHECK: (type $7 (func (param i32) (result f64)))

;; CHECK: (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
(type $struct.packed (struct (field (mut i8)) (field (mut i32))))

(type $struct.nondefaultable (struct (field (ref $struct.A))))

(type $struct.recursive (struct (field (mut (ref null $struct.recursive)))))

(type $struct.nonnullable (struct (field (ref $struct.A))))

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

;; CHECK: (type $9 (func (param (ref null $struct.recursive))))

;; CHECK: (type $10 (func (param (ref $struct.A))))
Expand Down Expand Up @@ -175,19 +175,114 @@
)

;; CHECK: (func $packed (type $1)
;; CHECK-NEXT: (local $temp (ref $struct.packed))
;; CHECK-NEXT: (local $1 i32)
;; CHECK-NEXT: (local $2 i32)
;; CHECK-NEXT: (local $3 i32)
;; CHECK-NEXT: (local $4 i32)
;; CHECK-NEXT: (local $5 i32)
;; CHECK-NEXT: (local $6 i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (local.set $3
;; CHECK-NEXT: (i32.const 1337)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $4
;; CHECK-NEXT: (i32.const 1338)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $1
;; CHECK-NEXT: (i32.and
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: (i32.const 255)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (local.get $4)
;; CHECK-NEXT: )
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $1
;; CHECK-NEXT: (i32.and
;; CHECK-NEXT: (i32.const 99998)
;; CHECK-NEXT: (i32.const 255)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (struct.get_u $struct.packed 0
;; CHECK-NEXT: (struct.new_default $struct.packed)
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $1)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (i32.const 99999)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (local.set $5
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $6
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $packed
;; We do not optimize packed structs yet.
(local $temp (ref $struct.packed))
;; Packed fields require masking of irrelevant bits, which we apply on the
;; sets.
(local.set $temp
(struct.new $struct.packed
(i32.const 1337)
(i32.const 1338)
)
)
(struct.set $struct.packed 0
(local.get $temp)
(i32.const 99998)
)
(drop
(struct.get $struct.packed 0
(struct.new_default $struct.packed)
(local.get $temp)
)
)
;; Unpacked fields in the same struct do not need anything.
(struct.set $struct.packed 1
(local.get $temp)
(i32.const 99999)
)
(drop
(struct.get $struct.packed 1
(local.get $temp)
)
)
;; When using struct.new_default we do not need any masking, as the values
;; written are 0 anyhow.
(local.set $temp
(struct.new_default $struct.packed)
)
)

;; CHECK: (func $with-init-values (type $1)
Expand Down Expand Up @@ -647,7 +742,7 @@
)
)

;; CHECK: (func $local-copies-conditional (type $8) (param $x i32) (result f64)
;; CHECK: (func $local-copies-conditional (type $7) (param $x i32) (result f64)
;; CHECK-NEXT: (local $ref (ref null $struct.A))
;; CHECK-NEXT: (local $2 i32)
;; CHECK-NEXT: (local $3 f64)
Expand Down Expand Up @@ -741,7 +836,7 @@
)
)

;; CHECK: (func $non-exclusive-get (type $8) (param $x i32) (result f64)
;; CHECK: (func $non-exclusive-get (type $7) (param $x i32) (result f64)
;; CHECK-NEXT: (local $ref (ref null $struct.A))
;; CHECK-NEXT: (local.set $ref
;; CHECK-NEXT: (struct.new_default $struct.A)
Expand Down Expand Up @@ -3311,3 +3406,124 @@
)
)
)

;; Packed arrays.
(module
;; CHECK: (type $0 (func (result i32)))

;; CHECK: (type $array8 (array (mut i8)))
(type $array8 (array (mut i8)))

;; CHECK: (type $array16 (array (mut i16)))
(type $array16 (array (mut i16)))

;; CHECK: (func $array8 (type $0) (result i32)
;; CHECK-NEXT: (local $temp (ref $array8))
;; CHECK-NEXT: (local $1 i32)
;; CHECK-NEXT: (local $2 i32)
;; CHECK-NEXT: (local $3 i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (local.set $1
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $3
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (i32.and
;; CHECK-NEXT: (i32.const 1337)
;; CHECK-NEXT: (i32.const 255)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $array8 (result i32)
(local $temp (ref $array8))
(local.set $temp
(array.new_default $array8
(i32.const 3)
)
)
(array.set $array8
(local.get $temp)
(i32.const 1)
(i32.const 1337)
)
(array.get $array8
(local.get $temp)
(i32.const 1)
)
)

;; CHECK: (func $array16 (type $0) (result i32)
;; CHECK-NEXT: (local $temp (ref $array16))
;; CHECK-NEXT: (local $1 i32)
;; CHECK-NEXT: (local $2 i32)
;; CHECK-NEXT: (local $3 i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result nullref)
;; CHECK-NEXT: (local.set $1
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $3
;; CHECK-NEXT: (i32.const 0)
;; CHECK-NEXT: )
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (i32.and
;; CHECK-NEXT: (i32.const 1337)
;; CHECK-NEXT: (i32.const 65535)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (ref.null none)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $array16 (result i32)
(local $temp (ref $array16))
(local.set $temp
(array.new_default $array16
(i32.const 3)
)
)
(array.set $array16
(local.get $temp)
(i32.const 1)
(i32.const 1337)
)
(array.get $array16
(local.get $temp)
(i32.const 1)
)
)
)
Loading

0 comments on commit 738e8fc

Please sign in to comment.