tensorgrad 0.0.1

package/dist/buffers.js

@@ -0,0 +1,114 @@
+// Buffer planning: walk a Graph and decide which GPU buffer each Tensor maps to.
+//
+// v1 strategy: one GPU buffer per IR Tensor. Static shapes mean every buffer's
+// size is known at compile time and lifetimes don't overlap between steps —
+// so no pooling needed. Total memory is the sum of every intermediate tensor.
+// For our transformer at B=256: ~30 MB of activations + grads. Easily fits.
+//
+// Categorization is what the runtime cares about:
+//   * param        — uploaded by user via uploadParams; persistent across steps
+//   * param_grad   — written each step by the backward pass; readable for inspection
+//   * tensor_input — uploaded each step (tokens, targets, masks)
+//   * intermediate — produced by an op; lifetime = within a single step
+//   * output       — special intermediate that should be made readable (loss)
+const dtypeBytes = { f32: 4, i32: 4, bool: 4 };
+function shapeSize(shape) {
+    let n = 1;
+    for (const d of shape)
+        n *= d;
+    return n;
+}
+/**
+ * Build a BufferPlan from a graph + the param-grad map produced by appendGrad.
+ * @param graph the full graph (forward + backward + any optimizer ops)
+ * @param paramGrads map from param name -> the Tensor that holds its gradient
+ * @param writebackDecls list of end-of-step writebacks (e.g. from appendAdam).
+ *                       Empty when there's no optimizer in the graph.
+ */
+export function planBuffers(graph, paramGrads, writebackDecls = []) {
+    const buffers = [];
+    const tensorToBuffer = new Map();
+    const paramsByName = new Map();
+    const inputsByName = new Map();
+    const paramGradsByName = new Map();
+    const statesByName = new Map();
+    // Build a quick reverse map: tensorId -> param name (for grads).
+    const gradTensorIdToName = new Map();
+    for (const [name, tensor] of Object.entries(paramGrads)) {
+        gradTensorIdToName.set(tensor.id, name);
+    }
+    // ...and tensorId -> param/input op (so we can name the buffer correctly).
+    const opByOutId = new Map();
+    for (const op of graph.ops)
+        opByOutId.set(op.out, op);
+    const outputSet = new Set(graph.outputs);
+    // Walk all tensors in id order. Categorize each.
+    for (const t of graph.tensors) {
+        const op = opByOutId.get(t.id);
+        let kind = 'intermediate';
+        let name = null;
+        let initValue;
+        if (op?.kind === 'param_input') {
+            kind = 'param';
+            name = op.name;
+        }
+        else if (op?.kind === 'tensor_input') {
+            kind = 'tensor_input';
+            name = op.name;
+        }
+        else if (op?.kind === 'state_input') {
+            kind = 'state';
+            name = op.name;
+            initValue = op.initValue;
+        }
+        else if (gradTensorIdToName.has(t.id)) {
+            kind = 'param_grad';
+            name = gradTensorIdToName.get(t.id);
+        }
+        else if (outputSet.has(t.id)) {
+            kind = 'output';
+        }
+        const spec = {
+            id: t.id,
+            byteSize: Math.max(4, shapeSize(t.shape) * dtypeBytes[t.dtype]),
+            dtype: t.dtype,
+            shape: t.shape,
+            kind,
+            name,
+            ...(initValue !== undefined ? { initValue } : {}),
+        };
+        buffers.push(spec);
+        tensorToBuffer.set(t.id, t.id); // 1:1 for v1
+        if (kind === 'param')
+            paramsByName.set(name, t.id);
+        if (kind === 'tensor_input')
+            inputsByName.set(name, t.id);
+        if (kind === 'param_grad')
+            paramGradsByName.set(name, t.id);
+        if (kind === 'state')
+            statesByName.set(name, t.id);
+    }
+    const outputBufferIds = graph.outputs.map(id => tensorToBuffer.get(id));
+    // Resolve writeback declarations to (source, dest) buffer-id pairs.
+    const writebacks = writebackDecls.map(decl => {
+        const sourceBufId = tensorToBuffer.get(decl.source.id);
+        if (sourceBufId === undefined) {
+            throw new Error(`planBuffers: writeback source tensor #${decl.source.id} not in graph`);
+        }
+        const destBufId = decl.destKind === 'param'
+            ? paramsByName.get(decl.destName)
+            : statesByName.get(decl.destName);
+        if (destBufId === undefined) {
+            throw new Error(`planBuffers: writeback dest ${decl.destKind}:'${decl.destName}' not found`);
+        }
+        const sourceSpec = buffers[sourceBufId];
+        const destSpec = buffers[destBufId];
+        if (sourceSpec.byteSize !== destSpec.byteSize) {
+            throw new Error(`planBuffers: writeback size mismatch for ${decl.destKind}:'${decl.destName}' ` +
+                `(source ${sourceSpec.byteSize} bytes vs dest ${destSpec.byteSize})`);
+        }
+        return { source: sourceBufId, dest: destBufId, bytes: sourceSpec.byteSize };
+    });
+    return { buffers, tensorToBuffer, paramsByName, inputsByName, paramGradsByName, statesByName, outputBufferIds, writebacks };
+}
+//# sourceMappingURL=buffers.js.map

package/dist/buffers.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"buffers.js","sourceRoot":"","sources":["../src/buffers.ts"],"names":[],"mappings":"AAAA,iFAAiF;AACjF,EAAE;AACF,+EAA+E;AAC/E,4EAA4E;AAC5E,8EAA8E;AAC9E,4EAA4E;AAC5E,EAAE;AACF,kDAAkD;AAClD,gFAAgF;AAChF,qFAAqF;AACrF,iEAAiE;AACjE,wEAAwE;AACxE,8EAA8E;AA0C9E,MAAM,UAAU,GAA0B,EAAE,GAAG,EAAE,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,IAAI,EAAE,CAAC,EAAE,CAAA;AAErE,SAAS,SAAS,CAAC,KAAY;IAC7B,IAAI,CAAC,GAAG,CAAC,CAAA;IACT,KAAK,MAAM,CAAC,IAAI,KAAK;QAAE,CAAC,IAAI,CAAC,CAAA;IAC7B,OAAO,CAAC,CAAA;AACV,CAAC;AAcD;;;;;;GAMG;AACH,MAAM,UAAU,WAAW,CACzB,KAAY,EACZ,UAAkC,EAClC,iBAAkC,EAAE;IAEpC,MAAM,OAAO,GAAiB,EAAE,CAAA;IAChC,MAAM,cAAc,GAAG,IAAI,GAAG,EAAkB,CAAA;IAChD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC9C,MAAM,gBAAgB,GAAG,IAAI,GAAG,EAAkB,CAAA;IAClD,MAAM,YAAY,GAAG,IAAI,GAAG,EAAkB,CAAA;IAE9C,iEAAiE;IACjE,MAAM,kBAAkB,GAAG,IAAI,GAAG,EAAkB,CAAA;IACpD,KAAK,MAAM,CAAC,IAAI,EAAE,MAAM,CAAC,IAAI,MAAM,CAAC,OAAO,CAAC,UAAU,CAAC,EAAE,CAAC;QACxD,kBAAkB,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,EAAE,IAAI,CAAC,CAAA;IACzC,CAAC;IACD,2EAA2E;IAC3E,MAAM,SAAS,GAAG,IAAI,GAAG,EAAkB,CAAA;IAC3C,KAAK,MAAM,EAAE,IAAI,KAAK,CAAC,GAAG;QAAE,SAAS,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,EAAE,EAAE,CAAC,CAAA;IAErD,MAAM,SAAS,GAAG,IAAI,GAAG,CAAC,KAAK,CAAC,OAAO,CAAC,CAAA;IAExC,iDAAiD;IACjD,KAAK,MAAM,CAAC,IAAI,KAAK,CAAC,OAAO,EAAE,CAAC;QAC9B,MAAM,EAAE,GAAG,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAA;QAC9B,IAAI,IAAI,GAAuB,cAAc,CAAA;QAC7C,IAAI,IAAI,GAAkB,IAAI,CAAA;QAC9B,IAAI,SAA6B,CAAA;QAEjC,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YAC/B,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,cAAc,EAAE,CAAC;YACvC,IAAI,GAAG,cAAc,CAAA;YACrB,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;QAChB,CAAC;aAAM,IAAI,EAAE,EAAE,IAAI,KAAK,aAAa,EAAE,CAAC;YACtC,IAAI,GAAG,OAAO,CAAA;YACd,IAAI,GAAG,EAAE,CAAC,IAAI,CAAA;YACd,SAAS,GAAG,EAAE,CAAC,SAAS,CAAA;QAC1B,CAAC;aAAM,IAAI,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YACxC,IAAI,GAAG,YAAY,CAAA;YACnB,IAAI,GAAG,kBAAkB,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAE,CAAA;QACtC,CAAC;aAAM,IAAI,SAAS,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC;YAC/B,IAAI,GAAG,QAAQ,CAAA;QACjB,CAAC;QAED,MAAM,IAAI,GAAe;YACvB,EAAE,EAAE,CAAC,CAAC,EAAE;YACR,QAAQ,EAAE,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC,KAAK,CAAC,GAAG,UAAU,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;YAC/D,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,KAAK,EAAE,CAAC,CAAC,KAAK;YACd,IAAI;YACJ,IAAI;YACJ,GAAG,CAAC,SAAS,KAAK,SAAS,CAAC,CAAC,CAAC,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SAClD,CAAA;QACD,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,CAAA;QAClB,cAAc,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA,CAAE,aAAa;QAE7C,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QACnD,IAAI,IAAI,KAAK,cAAc;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC1D,IAAI,IAAI,KAAK,YAAY;YAAE,gBAAgB,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;QAC5D,IAAI,IAAI,KAAK,OAAO;YAAE,YAAY,CAAC,GAAG,CAAC,IAAK,EAAE,CAAC,CAAC,EAAE,CAAC,CAAA;IACrD,CAAC;IAED,MAAM,eAAe,GAAG,KAAK,CAAC,OAAO,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,cAAc,CAAC,GAAG,CAAC,EAAE,CAAE,CAAC,CAAA;IAExE,oEAAoE;IACpE,MAAM,UAAU,GAAgB,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,EAAE;QACxD,MAAM,WAAW,GAAG,cAAc,CAAC,GAAG,CAAC,IAAI,CAAC,MAAM,CAAC,EAAE,CAAC,CAAA;QACtD,IAAI,WAAW,KAAK,SAAS,EAAE,CAAC;YAC9B,MAAM,IAAI,KAAK,CAAC,yCAAyC,IAAI,CAAC,MAAM,CAAC,EAAE,eAAe,CAAC,CAAA;QACzF,CAAC;QACD,MAAM,SAAS,GAAG,IAAI,CAAC,QAAQ,KAAK,OAAO;YACzC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC;YACjC,CAAC,CAAC,YAAY,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAA;QACnC,IAAI,SAAS,KAAK,SAAS,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,+BAA+B,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,aAAa,CAAC,CAAA;QAC9F,CAAC;QACD,MAAM,UAAU,GAAG,OAAO,CAAC,WAAW,CAAE,CAAA;QACxC,MAAM,QAAQ,GAAG,OAAO,CAAC,SAAS,CAAE,CAAA;QACpC,IAAI,UAAU,CAAC,QAAQ,KAAK,QAAQ,CAAC,QAAQ,EAAE,CAAC;YAC9C,MAAM,IAAI,KAAK,CACb,4CAA4C,IAAI,CAAC,QAAQ,KAAK,IAAI,CAAC,QAAQ,IAAI;gBAC/E,WAAW,UAAU,CAAC,QAAQ,kBAAkB,QAAQ,CAAC,QAAQ,GAAG,CACrE,CAAA;QACH,CAAC;QACD,OAAO,EAAE,MAAM,EAAE,WAAW,EAAE,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,UAAU,CAAC,QAAQ,EAAE,CAAA;IAC7E,CAAC,CAAC,CAAA;IAEF,OAAO,EAAE,OAAO,EAAE,cAAc,EAAE,YAAY,EAAE,YAAY,EAAE,gBAAgB,EAAE,YAAY,EAAE,eAAe,EAAE,UAAU,EAAE,CAAA;AAC7H,CAAC"}

package/dist/codegen.d.ts

@@ -0,0 +1,23 @@
+import type { Graph, OpNode } from './ir.js';
+import type { BufferPlan } from './buffers.js';
+export interface KernelSpec {
+    /** Index into graph.ops. */
+    opIndex: number;
+    /** Op kind (for debugging / pipeline cache key). */
+    opKind: OpNode['kind'];
+    /** Generated WGSL source. Empty string for "logical" ops with no kernel. */
+    wgsl: string;
+    /**
+     * Buffer ids in binding-index order. The runtime creates a bind group with
+     * these in @binding(0..N) on @group(0). Inputs come first (read), output last
+     * (read_write).
+     */
+    bindings: number[];
+    /** Number of threads to dispatch (1-D). 0 means "skip" (e.g. reshape no-op). */
+    threads: number;
+    /** Workgroup size; usually WG_SIZE. */
+    workgroupSize: number;
+}
+/** Generate a KernelSpec per compute op in graph.ops (in dispatch order). */
+export declare function emitKernels(graph: Graph, plan: BufferPlan): KernelSpec[];
+//# sourceMappingURL=codegen.d.ts.map

package/dist/codegen.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"codegen.d.ts","sourceRoot":"","sources":["../src/codegen.ts"],"names":[],"mappings":"AAYA,OAAO,KAAK,EAAE,KAAK,EAAE,MAAM,EAAiB,MAAM,SAAS,CAAA;AAC3D,OAAO,KAAK,EAAE,UAAU,EAAE,MAAM,cAAc,CAAA;AAO9C,MAAM,WAAW,UAAU;IACzB,4BAA4B;IAC5B,OAAO,EAAE,MAAM,CAAA;IACf,oDAAoD;IACpD,MAAM,EAAE,MAAM,CAAC,MAAM,CAAC,CAAA;IACtB,4EAA4E;IAC5E,IAAI,EAAE,MAAM,CAAA;IACZ;;;;OAIG;IACH,QAAQ,EAAE,MAAM,EAAE,CAAA;IAClB,gFAAgF;IAChF,OAAO,EAAE,MAAM,CAAA;IACf,uCAAuC;IACvC,aAAa,EAAE,MAAM,CAAA;CACtB;AAMD,6EAA6E;AAC7E,wBAAgB,WAAW,CAAC,KAAK,EAAE,KAAK,EAAE,IAAI,EAAE,UAAU,GAAG,UAAU,EAAE,CAQxE"}