diff --git a/src/passes/MemoryPacking.cpp b/src/passes/MemoryPacking.cpp index c475164851d..5ecbc587266 100644 --- a/src/passes/MemoryPacking.cpp +++ b/src/passes/MemoryPacking.cpp @@ -34,6 +34,7 @@ #include "ir/utils.h" #include "pass.h" #include "support/space.h" +#include "support/stdckdint.h" #include "wasm-binary.h" #include "wasm-builder.h" #include "wasm.h" @@ -46,6 +47,7 @@ namespace { // will be used instead of memory.init for this range. struct Range { bool isZero; + // The range [start, end) - that is, start is included while end is not. size_t start; size_t end; }; @@ -109,7 +111,8 @@ struct MemoryPacking : public Pass { ReferrersMap& referrers); bool canSplit(const std::unique_ptr& segment, const Referrers& referrers); - void calculateRanges(const std::unique_ptr& segment, + void calculateRanges(Module* module, + const std::unique_ptr& segment, const Referrers& referrers, std::vector& ranges); void createSplitSegments(Builder& builder, @@ -162,7 +165,7 @@ void MemoryPacking::run(Module* module) { std::vector ranges; if (canSplit(segment, currReferrers)) { - calculateRanges(segment, currReferrers, ranges); + calculateRanges(module, segment, currReferrers, ranges); } else { // A single range covers the entire segment. Set isZero to false so the // original memory.init will be used even if segment is all zeroes. @@ -295,7 +298,8 @@ bool MemoryPacking::canSplit(const std::unique_ptr& segment, return segment->isPassive || segment->offset->is(); } -void MemoryPacking::calculateRanges(const std::unique_ptr& segment, +void MemoryPacking::calculateRanges(Module* module, + const std::unique_ptr& segment, const Referrers& referrers, std::vector& ranges) { auto& data = segment->data; @@ -303,6 +307,24 @@ void MemoryPacking::calculateRanges(const std::unique_ptr& segment, return; } + // A segment that might trap during startup must be handled carefully, as we + // do not want to remove that trap (unless we are allowed to by TNH). + auto preserveTrap = !getPassOptions().trapsNeverHappen && segment->offset; + if (preserveTrap) { + // Check if we can rule out a trap by it being in bounds. + if (auto* c = segment->offset->dynCast()) { + auto* memory = module->getMemory(segment->memory); + auto memorySize = memory->initial * Memory::kPageSize; + Index start = c->value.getUnsigned(); + Index size = segment->data.size(); + Index end; + if (!std::ckd_add(&end, start, size) && end <= memorySize) { + // This is in bounds. + preserveTrap = false; + } + } + } + // Calculate initial zero and nonzero ranges size_t start = 0; while (start < data.size()) { @@ -346,7 +368,10 @@ void MemoryPacking::calculateRanges(const std::unique_ptr& segment, } } - // Merge edge zeroes if they are not worth splitting + // Merge edge zeroes if they are not worth splitting. Note that we must not + // have a trap we must preserve here (if we did, we'd need to handle that in + // the code below). + assert(!preserveTrap); if (ranges.size() >= 2) { auto last = ranges.end() - 1; auto penultimate = ranges.end() - 2; @@ -364,12 +389,12 @@ void MemoryPacking::calculateRanges(const std::unique_ptr& segment, } } } else { - // Legacy ballpark overhead of active segment with offset + // Legacy ballpark overhead of active segment with offset. // TODO: Tune this threshold = 8; } - // Merge ranges across small spans of zeroes + // Merge ranges across small spans of zeroes. std::vector mergedRanges = {ranges.front()}; size_t i; for (i = 1; i < ranges.size() - 1; ++i) { @@ -383,10 +408,28 @@ void MemoryPacking::calculateRanges(const std::unique_ptr& segment, mergedRanges.push_back(*curr); } } - // Add the final range if it hasn't already been merged in + // Add the final range if it hasn't already been merged in. if (i < ranges.size()) { mergedRanges.push_back(ranges.back()); } + // If we need to preserve a trap then we must keep the topmost byte of the + // segment, which is enough to cause a trap if we do in fact trap. + if (preserveTrap) { + // Check if the last byte is in a zero range. Such a range will be dropped + // later, so add a non-zero range with that byte. (It is slightly odd to + // add a range with a zero marked as non-zero, but that is how we ensure it + // is preserved later in the output.) + auto& back = mergedRanges.back(); + if (back.isZero) { + // Remove the last byte from |back|. Decrementing this prevents it from + // overlapping with the new segment we are about to add. Note that this + // might make |back| have size 0, but that is not a problem as it will be + // dropped later anyhow, since it contains zeroes. + back.end--; + auto lastByte = data.size() - 1; + mergedRanges.push_back({false, lastByte, lastByte + 1}); + } + } std::swap(ranges, mergedRanges); } @@ -551,6 +594,20 @@ void MemoryPacking::dropUnusedSegments( module->updateDataSegmentsMap(); } +// Given the start of a segment and an offset into it, compute the sum of the +// two to get the absolute address the data should be at. This takes into +// account overflows in that we saturate to UINT_MAX: we do not want an overflow +// to take us down into a small address; in the invalid case of an overflow we +// stay at the largest possible unsigned value, which will keep us trapping. +template +Expression* addStartAndOffset(T start, T offset, Builder& builder) { + T total; + if (std::ckd_add(&total, start, offset)) { + total = std::numeric_limits::max(); + } + return builder.makeConst(T(total)); +} + void MemoryPacking::createSplitSegments( Builder& builder, const DataSegment* segment, @@ -568,10 +625,12 @@ void MemoryPacking::createSplitSegments( if (!segment->isPassive) { if (auto* c = segment->offset->dynCast()) { if (c->value.type == Type::i32) { - offset = builder.makeConst(int32_t(c->value.geti32() + range.start)); + offset = addStartAndOffset( + range.start, c->value.geti32(), builder); } else { assert(c->value.type == Type::i64); - offset = builder.makeConst(int64_t(c->value.geti64() + range.start)); + offset = addStartAndOffset( + range.start, c->value.geti64(), builder); } } else { assert(ranges.size() == 1); diff --git a/test/lit/passes/memory-packing_traps.wast b/test/lit/passes/memory-packing_traps.wast new file mode 100644 index 00000000000..e7677ff845a --- /dev/null +++ b/test/lit/passes/memory-packing_traps.wast @@ -0,0 +1,85 @@ +;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited. + +;; RUN: foreach %s %t wasm-opt --memory-packing -all --zero-filled-memory -S -o - | filecheck %s +;; RUN: foreach %s %t wasm-opt --memory-packing -all --zero-filled-memory -tnh -S -o - | filecheck %s --check-prefix=TNH__ + +(module + ;; We should not optimize out a segment that will trap, as that is an effect + ;; we need to preserve (unless TrapsNeverHappen). + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + ;; CHECK: (data $data (i32.const -1) "\00") + (data $data (i32.const -1) "\00") +) + +(module + ;; We should handle the possible overflow in adding the offset and size, and + ;; see this might trap. To keep the segment trapping, we will emit a segment + ;; with offset -1 of size 1 (which is the minimal thing we need for a trap). + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + ;; CHECK: (data $data (i32.const -1) "\00") + (data $data (i32.const -2) "\00\00\00") +) + +(module + ;; This segment will almost trap, but not. + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + (data $data (i32.const 65535) "\00") +) + +(module + ;; This one is slightly larger, and will trap. We can at least shorten the + ;; segment to only contain one byte, at the highest address the segment would + ;; write to. + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + ;; CHECK: (data $data (i32.const 65536) "\00") + (data $data (i32.const 65535) "\00\00") +) + +(module + ;; This one is slightly larger, but the offset is lower so it will not trap. + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + (data $data (i32.const 65534) "\00\00") +) + +(module + ;; This one's offset is just large enough to trap. + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + ;; CHECK: (data $data (i32.const 65536) "\00") + (data $data (i32.const 65536) "\00") +) + +(module + ;; This offset is unknown, so assume the worst. + ;; TODO: We could remove it in TNH mode + + ;; CHECK: (import "a" "b" (global $g i32)) + ;; TNH__: (import "a" "b" (global $g i32)) + (import "a" "b" (global $g i32)) + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + ;; CHECK: (data $data (global.get $g) "\00") + ;; TNH__: (data $data (global.get $g) "\00") + (data $data (global.get $g) "\00") +) + +(module + ;; Passive segments cannot trap during startup and are removable if they have + ;; no uses, like here. + ;; CHECK: (memory $memory 1 2) + ;; TNH__: (memory $memory 1 2) + (memory $memory 1 2) + (data $data "\00\00\00") +)