This is an automated email from the ASF dual-hosted git repository.
akaashrp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 828d117ebd [Web] Expose RNG state for deterministic restore (#20034)
828d117ebd is described below
commit 828d117ebdb90e4474e5b7a9ead4e88b35865a58
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Jul 21 17:31:05 2026 -0700
[Web] Expose RNG state for deterministic restore (#20034)
Expose tvmjs RNG state so callers can save and restore deterministic
sampling progress. Adds the following:
- `RNGState` as the exported state type for
`LinearCongruentialGenerator`
- `LinearCongruentialGenerator.getState()`
- `LinearCongruentialGenerator.setState(state)`
- `Instance.getRNGState()`
- `Instance.setRNGState(state)`
---
web/src/index.ts | 1 +
web/src/runtime.ts | 21 ++++++++++++++++-
web/src/support.ts | 32 +++++++++++++++++++++----
web/tests/node/test_random_generator.js | 42 +++++++++++++++++++++++++++++++++
4 files changed, 91 insertions(+), 5 deletions(-)
diff --git a/web/src/index.ts b/web/src/index.ts
index 4de7e09303..7cd05e7621 100644
--- a/web/src/index.ts
+++ b/web/src/index.ts
@@ -40,6 +40,7 @@ export {
export { Disposable, LibraryProvider } from "./types";
export { RPCServer } from "./rpc_server";
export { assert, wasmPath, LinearCongruentialGenerator } from "./support";
+export type { RNGState } from "./support";
export { detectGPUDevice, GPUDeviceDetectOutput } from "./webgpu";
export { LRUCache, CacheState } from "./cache_state";
export { createPolyfillWASI } from "./compact";
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 078a0c7df2..9c83a4d9ac 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -23,7 +23,12 @@
import { Pointer, PtrOffset, SizeOf, TypeIndex } from "./ctypes";
import { Disposable } from "./types";
import { Memory, CachedCallStack } from "./memory";
-import { assert, StringToUint8Array, LinearCongruentialGenerator } from
"./support";
+import {
+ assert,
+ StringToUint8Array,
+ LinearCongruentialGenerator,
+ RNGState,
+} from "./support";
import { Environment } from "./environment";
import { AsyncifyHandler } from "./asyncify";
import { FunctionInfo, WebGPUContext } from "./webgpu";
@@ -1542,6 +1547,20 @@ export class Instance implements Disposable {
this.rng.setSeed(seed);
}
+ /**
+ * Get the state of the internal LinearCongruentialGenerator.
+ */
+ getRNGState(): RNGState {
+ return this.rng.getState();
+ }
+
+ /**
+ * Restore the state of the internal LinearCongruentialGenerator.
+ */
+ setRNGState(state: RNGState): void {
+ this.rng.setState(state);
+ }
+
/**
* Sample index via top-p sampling.
*
diff --git a/web/src/support.ts b/web/src/support.ts
index be85e85b7b..96cb1923b1 100644
--- a/web/src/support.ts
+++ b/web/src/support.ts
@@ -80,24 +80,26 @@ export function wasmPath(): string {
* Linear congruential generator for random number generating that can be
seeded.
*
* Follows the implementation of `include/tvm/support/random_engine.h`, which
follows the
- * sepcification in
https://en.cppreference.com/w/cpp/numeric/random/linear_congruential_engine.
+ * specification in
https://en.cppreference.com/w/cpp/numeric/random/linear_congruential_engine.
*
* Note `Number.MAX_SAFE_INTEGER = 2^53 - 1`, and our intermediates are
strictly less than 2^48.
*/
+export type RNGState = number;
+
export class LinearCongruentialGenerator {
readonly modulus: number;
readonly multiplier: number;
readonly increment: number;
- // Always within the range (0, 2^32 - 1) non-inclusive; if 0, will forever
generate 0.
+ // Always within the range (0, modulus) non-inclusive; if 0, will forever
generate 0.
private rand_state: number;
/**
* Set modulus, multiplier, and increment. Initialize `rand_state` according
to `Date.now()`.
*/
constructor() {
- this.modulus = 2147483647; // 2^32 - 1
- this.multiplier = 48271; // between 2^15 and 2^16
+ this.modulus = 2147483647; // 2^31 - 1
+ this.multiplier = 48271; // between 2^15 and 2^16
this.increment = 0;
this.setSeed(Date.now());
}
@@ -119,6 +121,28 @@ export class LinearCongruentialGenerator {
this.checkRandState();
}
+ /**
+ * Get the current generator state for deterministic restoration.
+ */
+ getState(): RNGState {
+ return this.rand_state;
+ }
+
+ /**
+ * Restore a state returned by `getState()`.
+ */
+ setState(state: RNGState): void {
+ if (!Number.isInteger(state)) {
+ throw new Error("RNG state should be an integer.");
+ }
+ if (state <= 0 || state >= this.modulus) {
+ throw new Error(
+ `RNG state should be an integer in (0, ${this.modulus}).`,
+ );
+ }
+ this.rand_state = state;
+ }
+
/**
* Generate the next integer in the range (0, this.modulus) non-inclusive,
updating `rand_state`.
*
diff --git a/web/tests/node/test_random_generator.js
b/web/tests/node/test_random_generator.js
index aefbf0f56f..0e8342202e 100644
--- a/web/tests/node/test_random_generator.js
+++ b/web/tests/node/test_random_generator.js
@@ -44,6 +44,28 @@ test("Test whether the same seed make two RNGs generate same
results", () => {
}
});
+test("Restoring RNG state reproduces next random floats", () => {
+ const rng1 = new tvmjs.LinearCongruentialGenerator();
+ const rng2 = new tvmjs.LinearCongruentialGenerator();
+ rng1.setSeed(42);
+ for (let i = 0; i < 8; i++) {
+ rng1.randomFloat();
+ }
+
+ const state = rng1.getState();
+ const expected = [];
+ for (let i = 0; i < 16; i++) {
+ expected.push(rng1.randomFloat());
+ }
+
+ rng2.setState(state);
+ const restored = [];
+ for (let i = 0; i < expected.length; i++) {
+ restored.push(rng2.randomFloat());
+ }
+ expect(restored).toEqual(expected);
+});
+
test("Test two RNGs with different seeds generate different results", () => {
const rng1 = new tvmjs.LinearCongruentialGenerator();
const rng2 = new tvmjs.LinearCongruentialGenerator();
@@ -67,3 +89,23 @@ test('Illegal argument to `setSeed()`', () => {
rng1.setSeed(42.5);
}).toThrow("Seed should be an integer.");
});
+
+test("Illegal argument to `setState()`", () => {
+ const rng = new tvmjs.LinearCongruentialGenerator();
+
+ expect(() => {
+ rng.setState(undefined);
+ }).toThrow("RNG state should be an integer.");
+ expect(() => {
+ rng.setState({});
+ }).toThrow("RNG state should be an integer.");
+ expect(() => {
+ rng.setState(0);
+ }).toThrow("RNG state should be an integer in");
+ expect(() => {
+ rng.setState(rng.modulus);
+ }).toThrow("RNG state should be an integer in");
+ expect(() => {
+ rng.setState(1.5);
+ }).toThrow("RNG state should be an integer.");
+});
