Add parallel viriformat relabel command (bench 5725647)

Co-authored-by: Codex <codex@openai.com>

README.md

@@ -114,6 +114,27 @@ Once `uci` is sent, Heimdall will switch to UCI mode: colored output will be tur
 
 The mixed mode interface can be exited from by pressing either Ctrl+C, Ctrl+D (these also work for UCI) or Esc and then confirming when prompted (mixed mode only)
 
+#### Relabelling viriformat data
+
+The `relabel` subcommand replaces every move score in a
+[viriformat 3.0.0](https://docs.rs/viriformat/3.0.0/viriformat/) file with a fresh Heimdall search score. Scores remain signed 16-bit, white-relative centipawn values.
+
+```sh
+heimdall relabel --input=games.vf --output=relabeled.vf --depth=10 \
+  --nodes-soft=5000 --nodes-hard=1000000 --hash=1 --threads=8 --join
+```
+
+The available options are:
+
+- `--input` and `--output`: Required input path and output path/prefix.
+- `--depth`: Maximum search depth (default: 10).
+- `--nodes-soft` and `--nodes-hard`: Per-position soft and hard node limits (defaults: 5,000 and 1,000,000).
+- `--hash`: Transposition-table size in MiB per worker (default: 1).
+- `--threads`: Number of independent relabelling workers (default: 1).
+- `--chunk-size`: Number of games read before work is split across workers (default: 1,024).
+- `--skip` and `--limit`: Skip or process at most this many games; zero means no limit.
+- `--join`: Concatenate the worker shards into exactly `--output`, preserving input game order, then remove the shards. Without it, output remains in files named `OUTPUT.part-000`, `OUTPUT.part-001`, and so on.
+
 
 ### Built-in TUI
 

src/heimdall.nim

@@ -14,7 +14,7 @@
 import std/[os, math, times, atomics, parseopt, strutils, strformat, options, random]
 
 import heimdall/[moves, board, search, movegen, position, transpositions, eval]
-import heimdall/util/[magics, limits, tunables, book_augment, logs]
+import heimdall/util/[magics, limits, tunables, book_augment, logs, relabel as relabelUtil]
 import heimdall/uci/session
 
 
@@ -104,9 +104,15 @@ when isMainModule:
         limit         = 0
         skip          = 0
         rounds        = 1
+        # Parameters for viriformat relabelling
+        relabel       = false
+        relabelInput  = none(string)
+        relabelOutput = none(string)
+        relabelChunk  = 1024
+        relabelJoin   = false
 
     var runTUI = false
-    const subcommands = ["magics", "testonly", "bench", "spsa", "chonk", "tui"]
+    const subcommands = ["magics", "testonly", "bench", "spsa", "chonk", "relabel", "tui"]
     for kind, key, value in parser.getopt():
         case kind:
             of cmdArgument:
@@ -118,7 +124,7 @@ when isMainModule:
                     benchDepth = key.parseInt()
                     continue
 
-                let inSubCommand = bench or getParams or magicGen or testOnly or augment
+                let inSubCommand = bench or getParams or magicGen or testOnly or augment or relabel
 
                 if key in subcommands and inSubCommand:
                     stderr.writeLine(&"heimdall: error: '{prevSubCmd}' subcommand does not accept any arguments")
@@ -147,6 +153,9 @@ when isMainModule:
                     of "chonk":
                         # Hehe me make chonky book
                         augment = true
+                    of "relabel":
+                        runUCI = false
+                        relabel = true
                     of "tui":
                         runUCI = false
                         runTUI = true
@@ -165,7 +174,8 @@ when isMainModule:
                             benchSilent = true
                         else:
                             stderr.writeLine(&"heimdall: bench: error: unknown long option '{key}'")
-                if augment:
+                            quit(-1)
+                elif augment:
                     case key:
                         of "input":
                             inputBook = some(value)
@@ -207,6 +217,33 @@ when isMainModule:
                         else:
                             stderr.writeLine(&"heimdall: chonk: error: unknown long option '{key}'")
                             quit(-1)
+                elif relabel:
+                    case key:
+                        of "input":
+                            relabelInput = some(value)
+                        of "output":
+                            relabelOutput = some(value)
+                        of "nodes-soft":
+                            searcherNodes.soft = parseBiggestUInt(value)
+                        of "nodes-hard":
+                            searcherNodes.hard = parseBiggestUInt(value)
+                        of "hash":
+                            searcherHash = parseBiggestUInt(value)
+                        of "depth":
+                            searcherDepth = parseBiggestInt(value)
+                        of "threads":
+                            threads = parseInt(value)
+                        of "chunk-size":
+                            relabelChunk = parseInt(value)
+                        of "limit":
+                            limit = parseInt(value)
+                        of "skip":
+                            skip = parseInt(value)
+                        of "join":
+                            relabelJoin = true
+                        else:
+                            stderr.writeLine(&"heimdall: relabel: error: unknown long option '{key}'")
+                            quit(-1)
                 else:
                     stderr.writeLine(&"heimdall: error: unknown long option '{key}'")
                     quit(-1)
@@ -228,7 +265,25 @@ when isMainModule:
                     quit(-1)
             of cmdEnd:
                 break
-    if not magicGen and not augment:
+    if relabel:
+        if not relabelInput.isSome() or not relabelOutput.isSome():
+            stderr.writeLine("heimdall: relabel: error: --input and --output are required")
+            quit(-1)
+        try:
+            relabelViriformat(relabelInput.get(), relabelOutput.get(), RelabelConfig(
+                depth: searcherDepth,
+                nodes: searcherNodes,
+                hashMiB: searcherHash,
+                threads: threads,
+                chunkSize: relabelChunk,
+                skip: skip,
+                limit: limit,
+                join: relabelJoin
+            ))
+        except CatchableError:
+            stderr.writeLine(&"heimdall: relabel: error: {getCurrentExceptionMsg()}")
+            quit(-1)
+    elif not magicGen and not augment:
         if runTUI:
             when defined(windows):
                 stderr.writeLine("heimdall: the built-in TUI is disabled on Windows because termios.h is unavailable")

src/heimdall/util/relabel.nim

@@ -0,0 +1,282 @@
+# Copyright 2026 Mattia Giambirtone & All Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import std/[math, os, strformat, syncio, times]
+
+import heimdall/[board, eval, movegen, search, transpositions]
+import heimdall/util/[limits, tunables, viriformat]
+import heimdall/util/memory/aligned
+
+
+type
+    RelabelConfig* = object
+        depth*: int
+        nodes*: tuple[soft, hard: uint64]
+        hashMiB*: uint64
+        threads*: int
+        chunkSize*: int
+        skip*: int
+        limit*: int
+        join*: bool
+
+    RelabelWork = object
+        stop: bool
+        games: seq[ViriformatGame]
+
+    RelabelResult = object
+        games: int
+        positions: int
+        offset: int64
+        bytes: int64
+        error: string
+
+    RelabelRange = tuple[worker: int, offset, bytes: int64]
+
+    RelabelWorker = object
+        outputPath: string
+        config: RelabelConfig
+        inbox: Channel[RelabelWork]
+        outbox: Channel[RelabelResult]
+
+    RelabelThread = Thread[ptr RelabelWorker]
+
+
+func shardPath(output: string, worker: int): string =
+    &"{output}.part-{worker:03}"
+
+
+proc relabelGame(game: ViriformatGame, searcher: var SearchManager,
+                 ttable: ptr TranspositionTable): string =
+    var
+        board = newChessboard(@[game.initial.position.clone()])
+        scores = newSeq[int16](game.moves.len)
+
+    for i, entry in game.moves:
+        # Decode first so malformed input is rejected before doing an expensive
+        # search or emitting a partial game.
+        let move = board.positions[^1].parseViriformatMove(entry.move)
+
+        searcher.setBoard(board.positions)
+        searcher.histories.clear()
+        ttable[].init(1)
+        let variations = searcher.search(silent=true)
+        if variations.len == 0 or variations[0].moves[0] == nullMove():
+            raise newException(ValueError, &"search produced no move for position {board.toFEN()}")
+
+        var score = variations[0].score
+        if board.sideToMove == Black:
+            score = -score
+        scores[i] = int16(score)
+        board.doMove(move)
+
+    result = game.toViriformat(scores)
+
+
+proc workerMain(worker: ptr RelabelWorker) {.thread.} =
+    var ttable: ptr TranspositionTable
+    try:
+        ttable = allocHeapAligned(TranspositionTable, 64)
+        if ttable == nil:
+            raise newException(IOError, "failed to allocate a transposition table object")
+        ttable[] = newTranspositionTable(worker.config.hashMiB * 1024 * 1024)
+
+        var searcher = newSearchManager(@[startpos()], ttable, getDefaultParameters(),
+                                        evalState=newEvalState(verbose=false), normalizeScore=false)
+        searcher.limiter.addLimit(newDepthLimit(worker.config.depth))
+        searcher.limiter.addLimit(newNodeLimit(worker.config.nodes.soft, worker.config.nodes.hard))
+
+        var output = syncio.open(worker.outputPath, fmWrite)
+        defer: output.close()
+
+        # Readiness handshake: the coordinator does not start reading a large
+        # input until every worker has its search state and output shard ready.
+        worker.outbox.send(RelabelResult())
+
+        while true:
+            let work = worker.inbox.recv()
+            if work.stop:
+                break
+
+            var response = RelabelResult()
+            try:
+                response.offset = output.getFilePos()
+                for game in work.games:
+                    let encoded = game.relabelGame(searcher, ttable)
+                    output.write(encoded)
+                    inc(response.games)
+                    inc(response.positions, game.moves.len)
+                output.flushFile()
+                response.bytes = output.getFilePos() - response.offset
+            except CatchableError:
+                response.error = getCurrentExceptionMsg()
+            worker.outbox.send(response)
+    except CatchableError:
+        worker.outbox.send(RelabelResult(error: getCurrentExceptionMsg()))
+    finally:
+        if ttable != nil:
+            ttable[].destroy()
+            freeHeapAligned(ttable)
+
+
+proc joinShards(output: string, shardPaths: openArray[string], ranges: openArray[RelabelRange]) =
+    var joined = syncio.open(output, fmWrite)
+    defer: joined.close()
+    var buffer: array[1024 * 1024, char]
+    for part in ranges:
+        let path = shardPaths[part.worker]
+        var shard = syncio.open(path, fmRead)
+        shard.setFilePos(part.offset)
+        var remaining = part.bytes
+        while remaining > 0:
+            let count = shard.readBuffer(addr buffer[0], min(remaining, buffer.len.int64).int)
+            if count == 0:
+                raise newException(IOError, &"short read while joining relabel shard '{path}'")
+            if joined.writeBuffer(addr buffer[0], count) != count:
+                raise newException(IOError, &"short write while joining relabel shard '{path}'")
+            dec(remaining, count)
+        shard.close()
+    joined.flushFile()
+
+
+proc validate(config: RelabelConfig) =
+    if config.depth < 1:
+        raise newException(ValueError, "depth must be at least 1")
+    if config.nodes.soft < 1 or config.nodes.hard < 1:
+        raise newException(ValueError, "soft and hard node limits must be at least 1")
+    if config.nodes.soft > config.nodes.hard:
+        raise newException(ValueError, "soft node limit cannot exceed the hard node limit")
+    if config.hashMiB < 1:
+        raise newException(ValueError, "hash size must be at least 1 MiB per worker")
+    if config.threads notin 1..1024:
+        raise newException(ValueError, "threads must be in 1..1024")
+    if config.chunkSize < 1:
+        raise newException(ValueError, "chunk size must be at least 1 game")
+    if config.skip < 0 or config.limit < 0:
+        raise newException(ValueError, "skip and limit cannot be negative")
+
+
+proc relabelViriformat*(inputPath, outputPath: string, config: RelabelConfig) =
+    ## Streams viriformat games in bounded chunks, searches each recorded
+    ## position on independent workers, and writes one output shard per worker.
+    ## If `join` is enabled the shards are concatenated and removed afterwards.
+    config.validate()
+    let inputAbsolute = inputPath.absolutePath()
+    if inputAbsolute == outputPath.absolutePath():
+        raise newException(ValueError, "input and output paths must differ")
+    for i in 0..<config.threads:
+        if inputAbsolute == shardPath(outputPath, i).absolutePath():
+            raise newException(ValueError, "input path collides with a generated output shard")
+
+    var input = syncio.open(inputPath, fmRead)
+    defer: input.close()
+
+    var
+        workers = newSeq[RelabelWorker](config.threads)
+        threads = newSeq[RelabelThread](config.threads)
+        paths = newSeq[string](config.threads)
+        startedWorkers = 0
+        openedWorkers = 0
+        workersStopped = false
+
+    defer:
+        if not workersStopped:
+            for i in 0..<startedWorkers:
+                workers[i].inbox.send(RelabelWork(stop: true))
+            for i in 0..<startedWorkers:
+                threads[i].joinThread()
+        for i in 0..<openedWorkers:
+            workers[i].inbox.close()
+            workers[i].outbox.close()
+
+    for i in 0..<config.threads:
+        paths[i] = shardPath(outputPath, i)
+        workers[i] = RelabelWorker(outputPath: paths[i], config: config)
+        workers[i].inbox.open()
+        workers[i].outbox.open()
+        inc(openedWorkers)
+        # Starting sequentially also avoids having all workers initialise the
+        # shared, immutable NNUE network at the same time.
+        threads[i].createThread(workerMain, addr workers[i])
+        inc(startedWorkers)
+        let ready = workers[i].outbox.recv()
+        if ready.error.len > 0:
+            raise newException(IOError, &"worker {i} failed to initialise: {ready.error}")
+
+    echo &"Relabelling '{inputPath}' with {config.threads} worker(s), depth={config.depth}, " &
+         &"nodes={config.nodes.soft}/{config.nodes.hard}, hash={config.hashMiB} MiB/worker"
+
+    var scratch: ViriformatGame
+    for _ in 0..<config.skip:
+        if not input.readViriformatGame(scratch):
+            break
+
+    let started = epochTime()
+    var
+        totalGames = 0
+        totalPositions = 0
+        eof = false
+        ranges: seq[RelabelRange]
+
+    while not eof and (config.limit == 0 or totalGames < config.limit):
+        let wanted = if config.limit == 0:
+                config.chunkSize
+            else:
+                min(config.chunkSize, config.limit - totalGames)
+        var
+            workerChunks = newSeq[seq[ViriformatGame]](config.threads)
+            chunkLen = 0
+            perWorker = wanted.ceilDiv(config.threads)
+        while chunkLen < wanted:
+            var game: ViriformatGame
+            if not input.readViriformatGame(game):
+                eof = true
+                break
+            workerChunks[min(chunkLen div perWorker, config.threads - 1)].add(game)
+            inc(chunkLen)
+        if chunkLen == 0:
+            break
+
+        for i in 0..<config.threads:
+            workers[i].inbox.send(RelabelWork(games: move(workerChunks[i])))
+
+        for i in 0..<config.threads:
+            let response = workers[i].outbox.recv()
+            if response.error.len > 0:
+                raise newException(ValueError, &"worker {i} failed: {response.error}")
+            inc(totalGames, response.games)
+            inc(totalPositions, response.positions)
+            if response.bytes > 0:
+                ranges.add((worker: i, offset: response.offset, bytes: response.bytes))
+
+        let
+            elapsed = epochTime() - started
+            rate = if elapsed > 0: totalPositions.float / elapsed else: 0.0
+        echo &"Relabelled {totalGames} games / {totalPositions} positions ({rate:.1f} positions/s)"
+
+    for worker in workers.mitems():
+        worker.inbox.send(RelabelWork(stop: true))
+    for thread in threads.mitems():
+        thread.joinThread()
+    workersStopped = true
+
+    if config.join:
+        echo &"Joining {paths.len} shards into '{outputPath}'"
+        joinShards(outputPath, paths, ranges)
+        for path in paths:
+            removeFile(path)
+    else:
+        echo &"Wrote {paths.len} shards: {outputPath}.part-000 ... {paths[^1]}"
+
+    let elapsed = epochTime() - started
+    echo &"Finished: {totalGames} games / {totalPositions} positions in {elapsed:.2f} seconds"

src/heimdall/util/viriformat.nim

@@ -0,0 +1,139 @@
+# Copyright 2026 Mattia Giambirtone & All Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+## Reader, writer and move decoder for viriformat 3.0.0.
+##
+## A game is a marlinformat position followed by (move, score) pairs and a
+## four-byte zero terminator. Viriformat numbers squares from a1 while
+## Heimdall numbers them from a8, so both move squares need a rank flip.
+
+import std/[endians, strformat, syncio]
+
+import heimdall/[board, movegen, moves, position]
+import heimdall/util/marlinformat
+
+
+const VIRIFORMAT_MOVE_SIZE* = 4
+
+
+type
+    ViriformatScoredMove* = object
+        move*: uint16
+        score*: int16
+
+    ViriformatGame* = object
+        ## `header` is kept verbatim so relabelling does not perturb metadata
+        ## that is unrelated to the move scores.
+        header*: string
+        initial*: MarlinFormatRecord
+        moves*: seq[ViriformatScoredMove]
+
+
+func readUint16LE(data: string, offset: int): uint16 {.inline.} =
+    littleEndian16(addr result, unsafeAddr data[offset])
+
+
+func readInt16LE(data: string, offset: int): int16 {.inline.} =
+    cast[int16](readUint16LE(data, offset))
+
+
+proc readExact(file: syncio.File, size: int, allowEof: bool): string =
+    result = newString(size)
+    let read = file.readBuffer(addr result[0], size)
+    if read == 0 and allowEof:
+        result.setLen(0)
+    elif read != size:
+        raise newException(IOError, &"truncated viriformat record: expected {size} bytes, read {read}")
+
+
+proc readViriformatGame*(file: syncio.File, game: var ViriformatGame): bool =
+    ## Reads one game. Returns false only for a clean EOF at a game boundary.
+    let header = file.readExact(RECORD_SIZE, allowEof=true)
+    if header.len == 0:
+        return false
+
+    game.header = header
+    game.initial = fromMarlinformat(header)
+    game.moves.setLen(0)
+
+    while true:
+        let entry = file.readExact(VIRIFORMAT_MOVE_SIZE, allowEof=false)
+        let
+            rawMove = readUint16LE(entry, 0)
+            score = readInt16LE(entry, 2)
+        if rawMove == 0 and score == 0:
+            break
+        if rawMove == 0:
+            raise newException(ValueError, "invalid viriformat move: the all-zero move encoding is reserved for the game terminator")
+        game.moves.add(ViriformatScoredMove(move: rawMove, score: score))
+    result = true
+
+
+func appendUint16LE(result: var string, value: uint16) {.inline.} =
+    var encoded: uint16
+    littleEndian16(addr encoded, unsafeAddr value)
+    for b in cast[array[2, char]](encoded):
+        result.add(b)
+
+
+func appendInt16LE(result: var string, value: int16) {.inline.} =
+    result.appendUint16LE(cast[uint16](value))
+
+
+func toViriformat*(game: ViriformatGame, scores: openArray[int16]): string =
+    ## Serialises a game with replacement move scores.
+    doAssert scores.len == game.moves.len
+    result = newStringOfCap(game.header.len + (game.moves.len + 1) * VIRIFORMAT_MOVE_SIZE)
+    result.add(game.header)
+    for i, entry in game.moves:
+        result.appendUint16LE(entry.move)
+        result.appendInt16LE(scores[i])
+    result.add("\0\0\0\0")
+
+
+proc parseViriformatMove*(position: var Position, encoded: uint16): Move =
+    ## Converts a viriformat move to Heimdall's move representation. Matching
+    ## against the generated legal moves supplies flags viriformat does not
+    ## encode explicitly, such as captures and double pawn pushes.
+    if encoded == 0:
+        raise newException(ValueError, "cannot decode viriformat's null/terminator move")
+
+    let
+        start = Square(encoded and 0x3f).flipRank()
+        target = Square((encoded shr 6) and 0x3f).flipRank()
+        promotion = int((encoded shr 12) and 0x3)
+        moveType = int(encoded shr 14)
+
+    var legalMoves = newMoveList()
+    position.generateMoves(legalMoves)
+    for candidate in legalMoves:
+        if candidate.startSquare != start or candidate.targetSquare != target:
+            continue
+
+        let matches = case moveType:
+            of 0:
+                not candidate.isPromotion() and not candidate.isEnPassant() and not candidate.isCastling()
+            of 1:
+                candidate.isEnPassant()
+            of 2:
+                candidate.isCastling()
+            of 3:
+                candidate.isPromotion() and candidate.flag().promotionToPiece() ==
+                    [Knight, Bishop, Rook, Queen][promotion]
+            else:
+                false
+        if matches:
+            return candidate
+
+    raise newException(ValueError, &"viriformat move 0x{encoded:04x} is illegal in position {position.toFEN()}")

tests/test_viriformat.nim

@@ -0,0 +1,59 @@
+import std/[os, syncio, tempfiles, unittest]
+
+import heimdall/[movegen, moves, pieces, position]
+import heimdall/util/[marlinformat, viriformat]
+
+
+func encodeMove(start, target: string, moveType: uint16 = 0,
+                promotion: uint16 = 0): uint16 =
+    let
+        source = start.toSquare().flipRank().uint16
+        destination = target.toSquare().flipRank().uint16
+    source or (destination shl 6) or (promotion shl 12) or (moveType shl 14)
+
+
+suite "viriformat 3.0.0":
+    test "normal moves acquire Heimdall's implicit flags":
+        var position = startpos()
+        let move = position.parseViriformatMove(encodeMove("e2", "e4"))
+        check move.toUCI() == "e2e4"
+        check move.flag() == DoublePush
+
+    test "special move types and square conversion":
+        var ep = fromFEN("4k3/8/8/3pP3/8/8/8/4K3 w - d6 0 1")
+        let epMove = ep.parseViriformatMove(encodeMove("e5", "d6", moveType=1))
+        check epMove.flag() == EnPassant
+
+        var castle = fromFEN("4k3/8/8/8/8/8/8/4K2R w K - 0 1")
+        let castleMove = castle.parseViriformatMove(encodeMove("e1", "h1", moveType=2))
+        check castleMove.flag() == ShortCastling
+
+        var promotion = fromFEN("4k3/P7/8/8/8/8/8/4K3 w - - 0 1")
+        let promotionMove = promotion.parseViriformatMove(
+            encodeMove("a7", "a8", moveType=3, promotion=3)
+        )
+        check promotionMove.flag() == PromotionQueen
+
+    test "games round-trip with replacement little-endian scores":
+        let record = createMarlinFormatRecord(startpos(), White, 0)
+        let game = ViriformatGame(
+            header: record.toMarlinformat(),
+            moves: @[
+                ViriformatScoredMove(move: encodeMove("e2", "e4"), score: 10),
+                ViriformatScoredMove(move: encodeMove("e7", "e5"), score: -20)
+            ]
+        )
+        let encoded = game.toViriformat([300'i16, -400'i16])
+        let (file, path) = createTempFile("heimdall-viriformat-", ".vf")
+        file.write(encoded)
+        file.setFilePos(0)
+        var decoded: ViriformatGame
+        check file.readViriformatGame(decoded)
+        check decoded.header == game.header
+        check decoded.moves.len == 2
+        check decoded.moves[0].move == game.moves[0].move
+        check decoded.moves[0].score == 300
+        check decoded.moves[1].score == -400
+        check not file.readViriformatGame(decoded)
+        file.close()
+        removeFile(path)