Q values implementation for Metal backend

cpp/command/analysis.cpp

@@ -35,6 +35,7 @@ struct AnalyzeRequest {
   bool includeMovesOwnershipStdev;
   bool includePolicy;
   bool includePVVisits;
+  bool includeQValues;
 
   bool reportDuringSearch;
   double reportDuringSearchEvery;
@@ -238,6 +239,7 @@ int MainCmds::analysis(const vector<string>& args) {
     "includeOwnershipStdev",
     "includePolicy",
     "includePVVisits",
+    "includeQValues",
     "reportDuringSearchEvery",
     "firstReportDuringSearchAfter",
     "priority",
@@ -323,6 +325,7 @@ int MainCmds::analysis(const vector<string>& args) {
       request->includeOwnership,request->includeOwnershipStdev,
       request->includeMovesOwnership,request->includeMovesOwnershipStdev,
       request->includePVVisits,
+      request->includeQValues,
       ret
     );
 
@@ -610,6 +613,7 @@ int MainCmds::analysis(const vector<string>& args) {
       rbase.includeMovesOwnershipStdev = false;
       rbase.includePolicy = false;
       rbase.includePVVisits = false;
+      rbase.includeQValues = false;
       rbase.reportDuringSearch = false;
       rbase.reportDuringSearchEvery = 1e30;
       rbase.firstReportDuringSearchAfter = 1e30;
@@ -1019,6 +1023,11 @@ int MainCmds::analysis(const vector<string>& args) {
         if(!suc)
           continue;
       }
+      if(input.find("includeQValues") != input.end()) {
+        bool suc = parseBoolean(input, "includeQValues", rbase.includeQValues, "Must be a boolean");
+        if(!suc)
+          continue;
+      }
       if(input.find("reportDuringSearchEvery") != input.end()) {
         bool suc = parseDouble(input, "reportDuringSearchEvery", rbase.reportDuringSearchEvery, 0.001, 1000000.0, "Must be number of seconds from 0.001 to 1000000.0");
         if(!suc)
@@ -1163,6 +1172,7 @@ int MainCmds::analysis(const vector<string>& args) {
           newRequest->includeMovesOwnershipStdev = rbase.includeMovesOwnershipStdev;
           newRequest->includePolicy = rbase.includePolicy;
           newRequest->includePVVisits = rbase.includePVVisits;
+          newRequest->includeQValues = rbase.includeQValues;
           newRequest->reportDuringSearch = rbase.reportDuringSearch;
           newRequest->reportDuringSearchEvery = rbase.reportDuringSearchEvery;
           newRequest->firstReportDuringSearchAfter = rbase.firstReportDuringSearchAfter;

cpp/command/contribute.cpp

@@ -261,7 +261,7 @@ static void runAndUploadSingleGame(
 
       // Usual analysis response fields
       ret["turnNumber"] = hist.moveHistory.size();
-      search->getAnalysisJson(perspective,analysisPVLen,preventEncore,true,alwaysIncludeOwnership,false,false,false,false,ret);
+      search->getAnalysisJson(perspective,analysisPVLen,preventEncore,true,alwaysIncludeOwnership,false,false,false,false,false,ret);
       std::cout << ret.dump() + "\n" << std::flush; // no endl due to race conditions
     }
 

cpp/command/evalsgf.cpp

@@ -641,6 +641,7 @@ int MainCmds::evalsgf(const vector<string>& args) {
       bool includeMovesOwnership = false;
       bool includeMovesOwnershipStdev = false;
       bool includePVVisits = true;
+      bool includeQValues = true;
       nlohmann::json ret;
       bool suc = search->getAnalysisJson(
         perspective,
@@ -652,6 +653,7 @@ int MainCmds::evalsgf(const vector<string>& args) {
         includeMovesOwnership,
         includeMovesOwnershipStdev,
         includePVVisits,
+        includeQValues,
         ret
       );
       if(suc) {

cpp/command/gtp.cpp

@@ -1656,6 +1656,47 @@ struct GTPEngine {
           out << endl;
         }
         out << endl;
+
+
+        if(nnOutput->whiteQWinloss != NULL) {
+          out << "whiteQWinloss" << endl;
+          for(int y = 0; y<board.y_size; y++) {
+            for(int x = 0; x<board.x_size; x++) {
+              int pos = NNPos::xyToPos(x,y,nnOutput->nnXLen);
+              if(nnOutput->policyProbs[pos] < 0)
+                out << "    NAN ";
+              else
+                out << Global::strprintf("%7.4f ", nnOutput->whiteQWinloss[pos]);
+            }
+            out << endl;
+          }
+          int pos = NNPos::locToPos(Board::PASS_LOC,board.x_size,nnOutput->nnXLen,nnOutput->nnYLen);
+          if(nnOutput->policyProbs[pos] < 0)
+            out << "    NAN ";
+          else
+            out << Global::strprintf("%7.4f ", nnOutput->whiteQWinloss[pos]);
+          out << endl;
+        }
+        if(nnOutput->whiteQScore != NULL) {
+          out << "whiteQScore" << endl;
+          for(int y = 0; y<board.y_size; y++) {
+            for(int x = 0; x<board.x_size; x++) {
+              int pos = NNPos::xyToPos(x,y,nnOutput->nnXLen);
+              if(nnOutput->policyProbs[pos] < 0)
+                out << "    NAN ";
+              else
+                out << Global::strprintf("%6.1f ", nnOutput->whiteQScore[pos]);
+            }
+            out << endl;
+          }
+          int pos = NNPos::locToPos(Board::PASS_LOC,board.x_size,nnOutput->nnXLen,nnOutput->nnYLen);
+          if(nnOutput->policyProbs[pos] < 0)
+            out << "    NAN ";
+          else
+            out << Global::strprintf("%6.1f ", nnOutput->whiteQScore[pos]);
+          out << endl;
+        }
+
       }
     }
 

cpp/neuralnet/cudabackend.cpp

@@ -2787,6 +2787,26 @@ void NeuralNet::getOutput(
       SymmetryHelpers::copyOutputsWithSymmetry(ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, inputBufs[row]->symmetry);
     }
 
+    if(output->whiteQWinloss != NULL || output->whiteQScore != NULL) {
+      testAssert(numPolicyChannels == 4 && modelVersion >= 16 && output->whiteQWinloss != NULL && output->whiteQScore != NULL);
+      if(gpuHandle->usingNHWC) {
+        for(int i = 0; i<nnXLen*nnYLen; i++)
+          policyProbsTmp[i] = policySrcBuf[i*numPolicyChannels+2];
+        SymmetryHelpers::copyOutputsWithSymmetry(policyProbsTmp, output->whiteQWinloss, 1, nnYLen, nnXLen, inputBufs[row]->symmetry);
+        for(int i = 0; i<nnXLen*nnYLen; i++)
+          policyProbsTmp[i] = policySrcBuf[i*numPolicyChannels+3];
+        SymmetryHelpers::copyOutputsWithSymmetry(policyProbsTmp, output->whiteQScore, 1, nnYLen, nnXLen, inputBufs[row]->symmetry);
+        output->whiteQWinloss[nnXLen*nnYLen] = policyPassSrcBuf[2];
+        output->whiteQScore[nnXLen*nnYLen] = policyPassSrcBuf[3];
+      }
+      else {
+        SymmetryHelpers::copyOutputsWithSymmetry(policySrcBuf+nnXLen*nnYLen*2, output->whiteQWinloss, 1, nnYLen, nnXLen, inputBufs[row]->symmetry);
+        SymmetryHelpers::copyOutputsWithSymmetry(policySrcBuf+nnXLen*nnYLen*3, output->whiteQScore, 1, nnYLen, nnXLen, inputBufs[row]->symmetry);
+        output->whiteQWinloss[nnXLen*nnYLen] = policyPassSrcBuf[2];
+        output->whiteQScore[nnXLen*nnYLen] = policyPassSrcBuf[3];
+      }
+    }
+
     if(modelVersion >= 9) {
       int numScoreValueChannels = gpuHandle->model->numScoreValueChannels;
       assert(numScoreValueChannels == 6);

cpp/neuralnet/metalbackend.cpp

@@ -797,6 +797,8 @@ void MetalProcess::processPolicy(
   const ComputeHandle* gpuHandle,
   NNResultBuf* inputBuf,
   size_t row) {
+  const int nnXLen = gpuHandle->nnXLen;
+  const int nnYLen = gpuHandle->nnYLen;
   auto& buffers = *inputBuffers;
   float* targetBuffer = &buffers.policyResults[row * buffers.singlePolicyResultElts * buffers.policyResultChannels];
   const auto symmetry = inputBuf->symmetry;
@@ -810,8 +812,19 @@ void MetalProcess::processPolicy(
     targetBuffer = &buffers.policyProbsBuffer[row * buffers.singlePolicyResultElts];
   }
 
-  SymmetryHelpers::copyOutputsWithSymmetry(
-    targetBuffer, currentOutput->policyProbs, 1, gpuHandle->nnYLen, gpuHandle->nnXLen, symmetry);
+  SymmetryHelpers::copyOutputsWithSymmetry(targetBuffer, currentOutput->policyProbs, 1, nnYLen, nnXLen, symmetry);
+
+  if(currentOutput->whiteQWinloss != NULL || currentOutput->whiteQScore != NULL) {
+    targetBuffer = &buffers.policyResults[row * buffers.singlePolicyResultElts * buffers.policyResultChannels];
+
+    SymmetryHelpers::copyOutputsWithSymmetry(
+      targetBuffer + (buffers.singlePolicyResultElts * 2), currentOutput->whiteQWinloss, 1, nnYLen, nnXLen, symmetry);
+    SymmetryHelpers::copyOutputsWithSymmetry(
+      targetBuffer + (buffers.singlePolicyResultElts * 3), currentOutput->whiteQScore, 1, nnYLen, nnXLen, symmetry);
+
+    currentOutput->whiteQWinloss[nnXLen * nnYLen] = buffers.policyPassResults[row * buffers.policyResultChannels + 2];
+    currentOutput->whiteQScore[nnXLen * nnYLen] = buffers.policyPassResults[row * buffers.policyResultChannels + 3];
+  }
 }
 
 void MetalProcess::processValue(

cpp/neuralnet/nneval.cpp

@@ -66,6 +66,7 @@ NNEvaluator::NNEvaluator(
   const vector<int>& gpuIdxByServerThr,
   const string& rSeed,
   bool doRandomize,
+  bool useQValues,
   int defaultSymmetry
 )
   :modelName(mName),
@@ -77,6 +78,7 @@ NNEvaluator::NNEvaluator(
    inputsUseNHWC(iUseNHWC),
    usingFP16Mode(useFP16Mode),
    usingNHWCMode(useNHWCMode),
+   usingQValues(useQValues),
    numThreads(numThr),
    gpuIdxByServerThread(gpuIdxByServerThr),
    randSeed(rSeed),
@@ -123,6 +125,9 @@ NNEvaluator::NNEvaluator(
       Global::intToString(nnXLen) + " * " + Global::intToString(nnYLen) +
       (requireExactNNLen ? " exactly" : " allowing smaller boards")
     );
+    logger->write(
+      "Enable raw NN Q values: " + Global::boolToString(usingQValues)
+    );
   }
 
   if(nnCacheSizePowerOfTwo >= 0)
@@ -518,6 +523,16 @@ void NNEvaluator::serve(
         else {
           resultBuf->result->whiteOwnerMap = NULL;
         }
+        if(usingQValues && modelVersion >= 16) {
+          float* whiteQWinloss = new float[nnXLen*nnYLen+1];
+          float* whiteQScore = new float[nnXLen*nnYLen+1];
+          for(int i = 0; i<nnXLen*nnYLen+1; i++) {
+            whiteQWinloss[i] = 0.0;
+            whiteQScore[i] = 0.0;
+          }
+          resultBuf->result->whiteQWinloss = whiteQWinloss;
+          resultBuf->result->whiteQScore = whiteQScore;
+        }
 
         //These aren't really probabilities. Win/Loss/NoResult will get softmaxed later
         double whiteWinProb = 0.0 + rand.nextGaussian() * 0.20;
@@ -552,6 +567,15 @@ void NNEvaluator::serve(
           emptyOutput->whiteOwnerMap = new float[nnXLen*nnYLen];
         else
           emptyOutput->whiteOwnerMap = NULL;
+
+        if(usingQValues && modelVersion >= 16) {
+          emptyOutput->whiteQWinloss = new float[nnXLen*nnYLen+1];
+          emptyOutput->whiteQScore = new float[nnXLen*nnYLen+1];
+        }
+        else {
+          emptyOutput->whiteQWinloss = NULL;
+          emptyOutput->whiteQScore = NULL;
+        }
         outputBuf.push_back(emptyOutput);
       }
 
@@ -852,6 +876,18 @@ void NNEvaluator::evaluate(
     buf.result->policyOptimismUsed = (float)resultWithoutOwnerMap->policyOptimismUsed;
     buf.result->nnXLen = resultWithoutOwnerMap->nnXLen;
     buf.result->nnYLen = resultWithoutOwnerMap->nnYLen;
+    if(resultWithoutOwnerMap->whiteQWinloss != NULL) {
+      assert(buf.result->whiteQWinloss != NULL);
+      assert(resultWithoutOwnerMap->nnXLen == nnXLen);
+      assert(resultWithoutOwnerMap->nnYLen == nnYLen);
+      std::copy(resultWithoutOwnerMap->whiteQWinloss, resultWithoutOwnerMap->whiteQWinloss + (nnXLen*nnYLen+1), buf.result->whiteQWinloss);
+    }
+    if(resultWithoutOwnerMap->whiteQScore != NULL) {
+      assert(buf.result->whiteQScore != NULL);
+      assert(resultWithoutOwnerMap->nnXLen == nnXLen);
+      assert(resultWithoutOwnerMap->nnYLen == nnYLen);
+      std::copy(resultWithoutOwnerMap->whiteQScore, resultWithoutOwnerMap->whiteQScore + (nnXLen*nnYLen+1), buf.result->whiteQScore);
+    }
     assert(buf.result->whiteOwnerMap != NULL);
   }
   else {
@@ -1126,6 +1162,31 @@ void NNEvaluator::evaluate(
     else {
       throw StringError("NNEval value postprocessing not implemented for model version");
     }
+
+    //Postprocess q values
+    if(buf.result->whiteQWinloss != NULL || buf.result->whiteQScore != NULL) {
+      if(!(buf.result->whiteQWinloss != NULL && buf.result->whiteQScore != NULL)) {
+        throw StringError("NNEval qvalue postprocessing found ony one of winloss and score");
+      }
+      for(int pos = 0; pos<nnXLen*nnYLen+1; pos++) {
+        if(!isLegal[pos]) {
+          buf.result->whiteQWinloss[pos] = 0.0f;
+          buf.result->whiteQScore[pos] = 0.0f;
+        }
+        else {
+          //Similarly as mentioned above, the result we get back from the net is actually not from white's perspective,
+          //but from the player to move, so we need to flip it to make it white at the same time as we tanh it.
+          if(nextPlayer == P_WHITE) {
+            buf.result->whiteQWinloss[pos] = tanh(buf.result->whiteQWinloss[pos]);
+            buf.result->whiteQScore[pos] = (float)(buf.result->whiteQScore[pos] * postProcessParams.scoreMeanMultiplier);
+          }
+          else {
+            buf.result->whiteQWinloss[pos] = -tanh(buf.result->whiteQWinloss[pos]);
+            buf.result->whiteQScore[pos] = -(float)(buf.result->whiteQScore[pos] * postProcessParams.scoreMeanMultiplier);
+          }
+        }
+      }
+    }
   }
 
   //Postprocess ownermap
@@ -1151,7 +1212,6 @@ void NNEvaluator::evaluate(
     }
   }
 
-
   //And record the nnHash in the result and put it into the table
   buf.result->nnHash = nnHash;
   if(nnCacheTable != NULL)

cpp/neuralnet/nneval.h

@@ -100,6 +100,7 @@ class NNEvaluator {
     const std::vector<int>& gpuIdxByServerThread,
     const std::string& randSeed,
     bool doRandomize,
+    bool useQValues,
     int defaultSymmetry
   );
   ~NNEvaluator();
@@ -217,6 +218,8 @@ class NNEvaluator {
   const bool inputsUseNHWC;
   const enabled_t usingFP16Mode;
   const enabled_t usingNHWCMode;
+  const bool usingQValues;
+
   int numThreads;
   std::vector<int> gpuIdxByServerThread;
   const std::string randSeed;