Humansl support in onnx trt path

Author: lightvector

cpp/neuralnet/onnxmodelbuilder.cpp

@@ -326,6 +326,21 @@ struct Builder {
     return addNode("Add", {maskName, c}, uniq(nameBase + "/gpmaskshift"), nameBase + "/gpmaskshift");
   }
 
+  // SGF metadata encoder (HumanSL nets, metaEncoderVersion > 0): a small MLP over the [N,metaC,1,1]
+  // metadata input producing a [N,trunkC,1,1] bias that gets added into the trunk's initial bias
+  // alongside the global-input matmul. Mirrors SGFMetadataEncoder::apply in eigenbackend.cpp and
+  // ModelParser::buildSGFMetadataEncoder in trtbackend.cpp. All ops are NC11 channel matmuls/biases,
+  // so this is layout-independent (it runs before any NCHW->NHWC trunk conversion).
+  string buildSGFMetadataEncoder(const string& input, const SGFMetadataEncoderDesc& desc) {
+    string x = buildMatMul(input, desc.mul1);
+    x = buildMatBias(x, desc.bias1);
+    x = buildActivation(x, desc.act1);
+    x = buildMatMul(x, desc.mul2);
+    x = buildMatBias(x, desc.bias2);
+    x = buildActivation(x, desc.act2);
+    return buildMatMul(x, desc.mul3);
+  }
+
   // ---- Residual block builders ----
   // useNHWC: input and output are channel-last [N,H,W,C], and the block's internals run NHWC. The
   // elementwise BN/activation/mask ops and 1x1 convs are layout-free; spatial convs (k>1) bubble to
@@ -814,9 +829,6 @@ Result build(
   bool transformerNHWC,
   Logger* logger
 ) {
-  if(desc.metaEncoderVersion > 0)
-    throw StringError("OnnxModelBuilder: SGF metadata encoder not yet supported");
-
   if(logger != NULL)
     logger->write("Building internal onnx model, requireExactNNLen=" + Global::boolToString(requireExactNNLen) + " transformerNHWC=" + Global::boolToString(transformerNHWC));
 
@@ -875,6 +887,12 @@ Result build(
   addInput("InputMask", 1);
   addInput("InputSpatial", numInputChannels);
   addInputNC11("InputGlobal", numInputGlobalChannels);
+  // HumanSL-style nets additionally take a per-row SGF metadata vector. Only declare the input when
+  // the model actually has an encoder - an unused graph input would just be dead weight (and the
+  // backend only allocates/binds an InputMeta buffer when numInputMetaChannels > 0).
+  bool hasMetaEncoder = desc.metaEncoderVersion > 0;
+  if(hasMetaEncoder)
+    addInputNC11("InputMeta", desc.numInputMetaChannels);
 
   // ---- Mask-derived features ----
   if(!requireExactNNLen) {
@@ -934,6 +952,12 @@ Result build(
   string initialConv = b.buildConv("InputSpatial", trunk.initialConv, false);
   string initialMatMul = b.buildMatMul("InputGlobal", trunk.initialMatMul);
   string cur = b.elementwise("Add", initialConv, initialMatMul, trunk.name + "/initbias");
+  if(hasMetaEncoder) {
+    testAssert(trunk.metaEncoderVersion > 0);
+    testAssert(trunk.sgfMetadataEncoder.mul3.outChannels == trunk.initialMatMul.outChannels);
+    string initialMeta = b.buildSGFMetadataEncoder("InputMeta", trunk.sgfMetadataEncoder);
+    cur = b.elementwise("Add", cur, initialMeta, trunk.name + "/initmetabias");
+  }
 
   // When transformerNHWC, run the entire trunk block stack channel-last: one NCHW->NHWC conversion
   // here and one NHWC->NCHW conversion before the trunk tip. Every block (convnet/gpool/nbt/

cpp/neuralnet/trtbackend.cpp

@@ -200,7 +200,10 @@ struct ModelParser {
 
   // Bump this when between katago versions we want to forcibly drop old timing caches and plan caches.
   // Bumped 7->8 for the TensorRT ONNX overhaul (ONNX emitter as default path, NHWC trunk, FP32 pinning).
-  static constexpr int tuneSalt = 8;
+  // Bumped 8->9 for SGF metadata encoder support on the ONNX path, and to discard caches potentially
+  // polluted by the concurrent-engine-build bug fixed in "Serialize TensorRT engine builds across GPU
+  // threads" (#1225).
+  static constexpr int tuneSalt = 9;
 
   unique_ptr<TRTModel> build(
     unique_ptr<INetworkDefinition> net,
@@ -1309,6 +1312,8 @@ struct ComputeHandle {
       setProfile("InputMask", Dims4(1, 1, ctx->nnYLen, ctx->nnXLen), Dims4(maxBatchSize, 1, ctx->nnYLen, ctx->nnXLen));
       setProfile("InputSpatial", Dims4(1, desc.numInputChannels, ctx->nnYLen, ctx->nnXLen), Dims4(maxBatchSize, desc.numInputChannels, ctx->nnYLen, ctx->nnXLen));
       setProfile("InputGlobal", Dims4(1, desc.numInputGlobalChannels, 1, 1), Dims4(maxBatchSize, desc.numInputGlobalChannels, 1, 1));
+      if(desc.metaEncoderVersion > 0)
+        setProfile("InputMeta", Dims4(1, desc.numInputMetaChannels, 1, 1), Dims4(maxBatchSize, desc.numInputMetaChannels, 1, 1));
 
       model = make_unique<TRTModel>();
       model->nnXLen = ctx->nnXLen;
@@ -1324,9 +1329,10 @@ struct ComputeHandle {
       // and the "nhwc" field distinguishes the NHWC vs NCHW trunk layout (different layer signatures),
       // so the two layouts don't share a timing-cache file full of mutual misses.
       string tuneDesc = Global::strprintf(
-        "\"onnxsalt\"(%d)\"nhwc\"(%d)\"model\"(%d,%d,%d)",
+        "\"onnxsalt\"(%d)\"nhwc\"(%d)\"model\"(%d,%d,%d,%d,%d)",
         ModelParser::tuneSalt, ctx->transformerNHWC ? 1 : 0,
-        desc.modelVersion, desc.numInputChannels, desc.numInputGlobalChannels);
+        desc.modelVersion, desc.numInputChannels, desc.numInputGlobalChannels,
+        desc.metaEncoderVersion, desc.numInputMetaChannels);
       SHA2::get256(tuneDesc.c_str(), model->tuneHash);
     }
     else {