@driftengine/texture 4.0.0

package/src/semantics.ts

@@ -0,0 +1,114 @@
+/**
+ * What a channel *is*, declared rather than conventional.
+ *
+ * **The family of defects this removes is the one nobody attributes correctly.** A normal map
+ * upside down in one asset and not in another; an albedo that is linear here and sRGB there; a
+ * gloss map read as roughness, which inverts every highlight in the scene. Each is a convention
+ * held in somebody's head, and each survives review because the file looks fine on its own.
+ *
+ * Declared, the format knows. A consumer never sees a convention at all, because `normaliseSample`
+ * has already applied it — a Y-down normal comes back Y-up, an sRGB value comes back linear, and
+ * gloss comes back as the roughness the engine shades with.
+ */
+export type ChannelSemantic =
+  | 'albedo-srgb'
+  | 'albedo-linear'
+  | 'normal-tangent-yup'
+  | 'normal-tangent-ydown'
+  | 'roughness-linear'
+  | 'gloss-linear'
+  | 'metallic-linear'
+  | 'occlusion-linear'
+  | 'height-linear'
+  | 'emissive-srgb'
+  | 'mask-linear';
+
+/**
+ * Every semantic, in the order a file stores them by.
+ *
+ * **A name crosses a package boundary as a number**, because `@driftengine/drft` is the container
+ * and knows nothing about what a channel means — a `DTEX` chunk stores `(semanticIndex << 4) |
+ * component` and would have to carry strings otherwise. The index is therefore part of the format:
+ * **a semantic may be appended and none may be reordered or removed**, or every file written before
+ * the change decodes its channels as something else. `semantics.test.ts` holds the order.
+ */
+export const CHANNEL_SEMANTICS: readonly ChannelSemantic[] = [
+  'albedo-srgb',
+  'albedo-linear',
+  'normal-tangent-yup',
+  'normal-tangent-ydown',
+  'roughness-linear',
+  'gloss-linear',
+  'metallic-linear',
+  'occlusion-linear',
+  'height-linear',
+  'emissive-srgb',
+  'mask-linear',
+];
+
+/** Where a semantic sits in that order, or −1 for one this build does not know. */
+export function semanticIndex(semantic: ChannelSemantic): number {
+  return CHANNEL_SEMANTICS.indexOf(semantic);
+}
+
+/** The semantic an index names, or null where a file names one from a later version. */
+export function semanticAt(index: number): ChannelSemantic | null {
+  return CHANNEL_SEMANTICS[index] ?? null;
+}
+
+export interface ChannelSpec {
+  readonly semantic: ChannelSemantic;
+  /** Which component of the decoded vector this channel occupies. */
+  readonly component: number;
+}
+
+/** Whether this channel carries colour, and therefore whether a transfer curve applies to it. */
+export function isColour(semantic: ChannelSemantic): boolean {
+  return semantic === 'albedo-srgb' || semantic === 'albedo-linear' || semantic === 'emissive-srgb';
+}
+
+/** Whether mipping this channel must preserve the normal distribution. See `mipNdf.ts`. */
+export function needsVarianceMips(semantic: ChannelSemantic): boolean {
+  return semantic === 'normal-tangent-yup' || semantic === 'normal-tangent-ydown';
+}
+
+/**
+ * The exact sRGB transfer function, piecewise, not the 2.2 power approximation.
+ *
+ * The approximation is within about two percent almost everywhere and wrong by more than that in
+ * the dark end, which in a texture pipeline is a colour shift nobody attributes to the right cause
+ * for weeks. The midpoint is the value that tells them apart: 0.5 linearises to 0.2140 exactly and
+ * to 0.2176 approximately.
+ */
+export function srgbToLinear(value: number): number {
+  return value <= 0.04045 ? value / 12.92 : Math.pow((value + 0.055) / 1.055, 2.4);
+}
+
+export function linearToSrgb(value: number): number {
+  return value <= 0.0031308 ? value * 12.92 : 1.055 * Math.pow(value, 1 / 2.4) - 0.055;
+}
+
+/**
+ * Write one raw channel value into its component of `out`, with its convention already applied.
+ *
+ * A consumer of this never asks which way a normal points or which curve an albedo carries.
+ */
+export function normaliseSample(out: Float32Array, spec: ChannelSpec, raw: number): void {
+  const at = spec.component;
+  switch (spec.semantic) {
+    case 'albedo-srgb':
+    case 'emissive-srgb':
+      out[at] = srgbToLinear(raw);
+      return;
+    case 'normal-tangent-ydown':
+      /* Flipped about the midpoint, because a tangent-space normal is stored biased into 0..1. */
+      out[at] = 1 - raw;
+      return;
+    case 'gloss-linear':
+      /* The engine shades with roughness. One of the two has to win, and it is not this one. */
+      out[at] = 1 - raw;
+      return;
+    default:
+      out[at] = raw;
+  }
+}

package/src/tensor/architecture.ts

@@ -0,0 +1,140 @@
+/**
+ * A network's definition, written as a function of its weights, built into a graph the runtime runs.
+ *
+ * **One definition, run over two sources.** At conversion it reads a checkpoint; afterwards it reads
+ * the converted file, whenever the runtime needs the graph at a new size — a vision transformer's
+ * patch grid follows the image's aspect ratio, and every device kernel bakes its shapes, so a model
+ * is one graph per size and not one graph. The two runs agree because of three rules:
+ *
+ * - **A weight is read by name, at the shape the definition expects**, and a name the source lacks
+ *   or a shape it disagrees with is refused naming both.
+ * - **A weight derived from others** — a batch norm folded into its convolution — is computed once,
+ *   at conversion, and kept under its own name; a source that already holds that name answers with
+ *   it and the computation never runs, which is why a converted file needs none of the parts.
+ * - **A constant** — a table computed from the shapes, such as rotary angles — carries a name
+ *   beginning `@`, is rebuilt on every run, and is never taken for a weight.
+ *
+ * **And every weight the source holds is read, derived from, or set aside by name.** A forgotten
+ * bias or layer scale is a network that runs and answers wrongly — the one failure nothing later can
+ * see — so a weight left over is refused. What that costs is an explicit `ignore` for every head a
+ * definition does not run, which is the point.
+ */
+import { validateGraph, type GraphTensor, type GraphValue, type NetworkGraph } from './graph.ts';
+import type { Attributes } from './operators.ts';
+
+/** Where a definition's weights come from: a checkpoint being converted, or a converted file. */
+export interface WeightSource {
+  get(name: string): GraphTensor | undefined;
+  names(): Iterable<string>;
+}
+
+/** What a constant's name begins with: a table of the shapes, never a weight. */
+export const CONSTANT_PREFIX = '@';
+
+export interface Weights {
+  /** The weight `name` as a value of the graph; refused if absent or held at another shape. */
+  read(name: string, shape?: readonly number[]): string;
+  /**
+   * A weight computed once from others, kept under `name`. `compute` reads its parts through
+   * `values`, which counts each as read; a source already holding `name` answers with it instead.
+   */
+  derive(
+    name: string,
+    shape: readonly number[],
+    compute: (values: (part: string, shape?: readonly number[]) => Float32Array) => Float32Array,
+  ): string;
+  /** A table computed from the graph's shapes, rebuilt on every run and never a weight. */
+  constant(name: string, shape: readonly number[], data: Float32Array): string;
+  /** Weights the definition deliberately does not read: a name, or a prefix ending in a dot. */
+  ignore(nameOrPrefix: string): void;
+}
+
+export interface GraphBuilder {
+  /** A node; its output is named `output` when given, and otherwise by the builder. */
+  node(op: string, inputs: readonly string[], attributes?: Attributes, output?: string): string;
+}
+
+/** A definition: it reads its weights, adds its nodes, and says what goes in and comes out. */
+export type Architecture = (
+  weights: Weights,
+  graph: GraphBuilder,
+) => { readonly inputs: readonly GraphValue[]; readonly outputs: readonly string[] };
+
+const shapeText = (shape: readonly number[]): string => `[${shape.join(', ')}]`;
+
+export function graphFromWeights(source: WeightSource, architecture: Architecture): NetworkGraph {
+  const tensors = new Map<string, GraphTensor>();
+  const read = new Set<string>();
+  const ignored: string[] = [];
+  const nodes: NetworkGraph['nodes'][number][] = [];
+  const take = (name: string, shape?: readonly number[]): GraphTensor => {
+    const tensor = source.get(name);
+    if (tensor === undefined) {
+      throw new Error(`the source has no weight "${name}", which the definition reads`);
+    }
+    if (shape !== undefined && shapeText(shape) !== shapeText(tensor.shape)) {
+      throw new Error(
+        `"${name}" is ${shapeText(tensor.shape)} in the source and the definition reads it as ` +
+          shapeText(shape),
+      );
+    }
+    read.add(name);
+    return tensor;
+  };
+  const weights: Weights = {
+    read(name, shape) {
+      if (!tensors.has(name)) tensors.set(name, take(name, shape));
+      else take(name, shape);
+      return name;
+    },
+    derive(name, shape, compute) {
+      if (tensors.has(name)) return name;
+      const held = source.get(name);
+      const data =
+        held === undefined
+          ? compute((part, partShape) => take(part, partShape).data)
+          : take(name, shape).data;
+      if (data.length !== shape.reduce((total, d) => total * d, 1)) {
+        throw new Error(
+          `the derived weight "${name}" has ${data.length} values for ${shapeText(shape)}`,
+        );
+      }
+      tensors.set(name, { shape, data });
+      return name;
+    },
+    constant(name, shape, data) {
+      const key = `${CONSTANT_PREFIX}${name}`;
+      tensors.set(key, { shape, data });
+      return key;
+    },
+    ignore(nameOrPrefix) {
+      ignored.push(nameOrPrefix);
+    },
+  };
+  let count = 0;
+  const graph: GraphBuilder = {
+    node(op, inputs, attributes = {}, output = `v${(count += 1)}`) {
+      nodes.push({ op, inputs, output, attributes });
+      return output;
+    },
+  };
+  const { inputs, outputs } = architecture(weights, graph);
+
+  const aside = (name: string): boolean =>
+    name.startsWith(CONSTANT_PREFIX) ||
+    ignored.some((rule) => (rule.endsWith('.') ? name.startsWith(rule) : name === rule));
+  const unread = [...source.names()].filter((name) => !read.has(name) && !aside(name));
+  if (unread.length > 0) {
+    const named = unread.slice(0, 5).map((name) => `"${name}"`);
+    throw new Error(
+      `the definition never reads ${named.join(', ')}` +
+        (unread.length > 5 ? ` and ${unread.length - 5} more` : '') +
+        ': a weight left behind is a network that runs and answers wrongly, so read it or set it ' +
+        'aside by name',
+    );
+  }
+  const network: NetworkGraph = { inputs, outputs, nodes, tensors };
+  const problem = validateGraph(network);
+  if (problem !== null) throw new Error(`the definition's graph cannot run: ${problem}`);
+  return network;
+}

package/src/tensor/attention.ts

@@ -0,0 +1,75 @@
+/**
+ * Multi-head attention, as a composition of the dense operators.
+ *
+ * **Heads are contiguous runs of channels, in the upstream order**: head `h` of `heads` owns
+ * channels `h·d` to `(h+1)·d − 1`, where `d` is the channel count over the head count. A port that
+ * interleaves them instead evaluates a network nobody trained.
+ *
+ * The queries, keys and values arrive already projected — the projections are matrix multiplies a
+ * graph states on its own — so this is the part that is attention's alone: scores scaled by
+ * `1/√d`, a stable softmax over each query's row, and the weighted sum of values.
+ *
+ * **A bias is added to the scaled scores**, `[heads][queries][keys]`, as a relative-position table
+ * or a mask is upstream; and **a batch is independent windows**, each attending over its own keys
+ * with the one bias. Queries and keys may differ in number, as they do where one set of tokens
+ * reads another.
+ */
+import { softmax } from './linear.ts';
+
+/**
+ * `out[batch][queries][channels]` from `q` of that shape and `k` and `v` of
+ * `[batch][keys][channels]`. `scratch` holds at least `queries × keys` values and is overwritten.
+ */
+export function attention(
+  out: Float32Array,
+  q: Float32Array,
+  k: Float32Array,
+  v: Float32Array,
+  queries: number,
+  keys: number,
+  channels: number,
+  heads: number,
+  scratch: Float32Array,
+  bias: Float32Array | null = null,
+  batch = 1,
+): void {
+  const d = channels / heads;
+  if (!Number.isInteger(d)) {
+    throw new RangeError(`attention: ${channels} channels do not split into ${heads} heads`);
+  }
+  const scale = 1 / Math.sqrt(d);
+  for (let b = 0; b < batch; b += 1) {
+    const qAt = b * queries * channels;
+    const kAt = b * keys * channels;
+    for (let head = 0; head < heads; head += 1) {
+      const first = head * d;
+      for (let query = 0; query < queries; query += 1) {
+        for (let key = 0; key < keys; key += 1) {
+          let dot = 0;
+          for (let c = 0; c < d; c += 1) {
+            dot +=
+              (q[qAt + query * channels + first + c] as number) *
+              (k[kAt + key * channels + first + c] as number);
+          }
+          const at = query * keys + key;
+          scratch[at] = dot * scale;
+          if (bias !== null) {
+            scratch[at] = (scratch[at] as number) + (bias[head * queries * keys + at] as number);
+          }
+        }
+      }
+      softmax(scratch, queries, keys, scratch);
+      for (let query = 0; query < queries; query += 1) {
+        for (let c = 0; c < d; c += 1) {
+          let sum = 0;
+          for (let key = 0; key < keys; key += 1) {
+            sum +=
+              (scratch[query * keys + key] as number) *
+              (v[kAt + key * channels + first + c] as number);
+          }
+          out[qAt + query * channels + first + c] = sum;
+        }
+      }
+    }
+  }
+}

package/src/tensor/denseOperators.ts

@@ -0,0 +1,153 @@
+/** The dense operators' rows: projections, activations, normalisation, attention. */
+import { attention } from './attention.ts';
+import { gelu, layerNorm, sigmoid, softmax } from './linear.ts';
+import { type Operator, elementwise, num, product, same, type Shape } from './operatorKit.ts';
+
+export const DENSE_OPERATORS: readonly (readonly [string, Operator])[] = [
+  [
+    'linear',
+    {
+      ranks: [2, 2, 1],
+      arity: [2, 3],
+      shape: ([x, w, b]) => {
+        const [rows, width] = x as Shape;
+        const [outs, ins] = w as Shape;
+        if (width !== ins) return `x has ${width} columns and the weight takes ${ins}`;
+        if (b !== undefined && b[0] !== outs) return `the bias has ${b[0]} and the weight ${outs}`;
+        return [rows as number, outs as number];
+      },
+      evaluate: ([x, w, b], [xs, ws], _attributes, out) => {
+        const [rows, ins] = xs as Shape as [number, number];
+        const outs = (ws as Shape)[0] as number;
+        for (let row = 0; row < rows; row += 1) {
+          for (let o = 0; o < outs; o += 1) {
+            let dot = 0;
+            for (let i = 0; i < ins; i += 1) {
+              dot +=
+                ((x as Float32Array)[row * ins + i] as number) *
+                ((w as Float32Array)[o * ins + i] as number);
+            }
+            /* The bias joins the double-precision sum, and the result is rounded once. */
+            out[row * outs + o] = dot + (b === undefined ? 0 : (b[o] as number));
+          }
+        }
+      },
+    },
+  ],
+  [
+    'relu',
+    {
+      arity: [1, 1],
+      shape: ([x]) => [...(x as Shape)],
+      evaluate: ([x], _shapes, _attributes, out) => {
+        for (let i = 0; i < out.length; i += 1)
+          out[i] = Math.max(0, (x as Float32Array)[i] as number);
+      },
+    },
+  ],
+  [
+    'gelu',
+    {
+      arity: [1, 1],
+      shape: ([x]) => [...(x as Shape)],
+      evaluate: ([x], _shapes, _attributes, out) => gelu(x as Float32Array, out),
+    },
+  ],
+  [
+    'sigmoid',
+    {
+      arity: [1, 1],
+      shape: ([x]) => [...(x as Shape)],
+      evaluate: ([x], _shapes, _attributes, out) => sigmoid(x as Float32Array, out),
+    },
+  ],
+  ['add', elementwise((a, b) => a + b)],
+  ['mul', elementwise((a, b) => a * b)],
+  [
+    'layerNorm',
+    {
+      ranks: [2, 1, 1],
+      arity: [3, 3],
+      shape: ([x, gamma]) => {
+        const cols = (x as Shape)[1];
+        return gamma?.[0] === cols ? [...(x as Shape)] : `gamma has ${gamma?.[0]} and x ${cols}`;
+      },
+      evaluate: ([x, gamma, beta], [xs], attributes, out) => {
+        const [rows, cols] = xs as Shape as [number, number];
+        const epsilon = num(attributes, 'epsilon', 1e-6);
+        layerNorm(
+          x as Float32Array,
+          rows,
+          cols,
+          gamma as Float32Array,
+          beta as Float32Array,
+          epsilon,
+          out,
+        );
+      },
+    },
+  ],
+  [
+    'softmax',
+    {
+      arity: [1, 1],
+      shape: ([x]) => [...(x as Shape)],
+      evaluate: ([x], [xs], _attributes, out) => {
+        const shape = xs as Shape;
+        const cols = shape[shape.length - 1] as number;
+        softmax(x as Float32Array, product(shape) / cols, cols, out);
+      },
+    },
+  ],
+  [
+    'attention',
+    {
+      /*
+       * Rank 2, or 3 for a batch of windows; checked here rather than by `ranks`, which states one.
+       * The bias, where there is one, is a score for each head, query and key.
+       */
+      arity: [3, 4],
+      shape: ([q, k, v, b], attributes) => {
+        const qs = q as Shape;
+        const ks = k as Shape;
+        if (qs.length !== 2 && qs.length !== 3)
+          return `the queries have rank ${qs.length}, not 2 or 3`;
+        if (ks.length !== qs.length)
+          return `the keys have rank ${ks.length} and the queries ${qs.length}`;
+        if (!same(ks, v as Shape)) return 'the keys and values differ in shape';
+        if (qs.length === 3 && qs[0] !== ks[0]) {
+          return `a batch of ${qs[0]} queries and of ${ks[0]} keys`;
+        }
+        const [queries, channels] = qs.slice(-2) as [number, number];
+        const [keys, keyChannels] = ks.slice(-2) as [number, number];
+        if (channels !== keyChannels) {
+          return `the queries have ${channels} channels and the keys ${keyChannels}`;
+        }
+        const heads = num(attributes, 'heads', 1);
+        if (channels % heads !== 0) return `${channels} channels do not split into ${heads} heads`;
+        if (b !== undefined && !same(b, [heads, queries, keys])) {
+          return `the bias is [${b.join(', ')}], not [${heads}, ${queries}, ${keys}]`;
+        }
+        return [...qs];
+      },
+      scratch: ([q, k]) => ((q as Shape).at(-2) as number) * ((k as Shape).at(-2) as number),
+      evaluate: ([q, k, v, b], [qs, ks], attributes, out, scratch) => {
+        const shape = qs as Shape;
+        const [queries, channels] = shape.slice(-2) as [number, number];
+        attention(
+          out,
+          q as Float32Array,
+          k as Float32Array,
+          v as Float32Array,
+          queries,
+          (ks as Shape).at(-2) as number,
+          channels,
+          num(attributes, 'heads', 1),
+          scratch,
+          b ?? null,
+          shape.length === 3 ? (shape[0] as number) : 1,
+        );
+      },
+    },
+  ],
+];