callBatch method

List<Lc0Output> callBatch(List<Tensor> inputs)

Batched forward: feed B positions in one call, get B outputs.

The conv tower runs as a single B*64, C matmul per layer, which is the whole reason batching helps: matmul cost scales close to O(B) instead of O(B) launches of a smaller kernel. Only the SE unit's gate is serial per batch element — the broadcast-by-matmul trick used at B=1 doesn't fold nicely across batch rows.

Implementation

List<Lc0Output> callBatch(List<Tensor> inputs) {
  if (inputs.isEmpty) return const <Lc0Output>[];
  final b = inputs.length;
  for (final t in inputs) {
    if (t.shape.length != 4 ||
        t.shape[0] != 1 ||
        t.shape[1] != 112 ||
        t.shape[2] != 8 ||
        t.shape[3] != 8) {
      throw ArgumentError(
        'Lc0Net.callBatch: each input must be [1, 112, 8, 8]; got ${t.shape}',
      );
    }
  }
  return Tensor.noGrad(() {
    final nhwc = _initialNHWCBatch(inputs); // [B*64, 112] on device
    var h = (_convForward(nhwc, inputConv, b) + inputConv.bias).relu();

    for (int i = 0; i < w.residual.length; i++) {
      h = _residualBatch(h, i, b);
    }

    final p1 = (_convForward(h, policy1, b) + policy1.bias).relu();
    final policyFlat = _convForward(p1, policyOut, b) + policyOut.bias;

    var v = (_convForward(h, valueConv, b) + valueConv.bias).relu();
    // v is [B*64, Vf]; row-major reshape to [B, 64*Vf] gives (s, c)
    // inner order — matches the pre-reordered ip1ValWT layout.
    v = v.reshape([b, w.valueFilters * 64]);
    v = (v.matmul(ip1ValWT) + ip1ValB).relu();
    v = v.matmul(ip2ValWT) + ip2ValB;
    if (w.wdl == 3) v = v.softmax();

    // Split per-batch on CPU — policy is [B*64, P], value is [B, wdl].
    final polHost = policyFlat.toFloat32List();
    final valHost = v.toFloat32List();
    final p = w.policyOutputPlanes;
    final results = <Lc0Output>[];
    for (int bi = 0; bi < b; bi++) {
      // Extract this batch's [64, P] slice, then convert to
      // [1, P, 8, 8] NCHW to match the caller's expectations.
      final polNCHW = Float32List(p * 64);
      for (int s = 0; s < 64; s++) {
        for (int pi = 0; pi < p; pi++) {
          polNCHW[pi * 64 + s] = polHost[(bi * 64 + s) * p + pi];
        }
      }
      final valSlice = Float32List(w.wdl);
      for (int k = 0; k < w.wdl; k++) {
        valSlice[k] = valHost[bi * w.wdl + k];
      }
      results.add(
        Lc0Output(
          Tensor.fromFloat32List([1, p, 8, 8], polNCHW, device: Device.CPU),
          Tensor.fromFloat32List([1, w.wdl], valSlice, device: Device.CPU),
        ),
      );
    }
    return results;
  });
}