Add uncertainty weighting to dag-preview

OpenBench/Makefile

@@ -21,4 +21,4 @@ endif
 all:
 	chmod +x ../build.sh
 	../build.sh $(BUILD_FLAGS) && mv ../build/release/lc0 $(EXE)
-	$(POST_BUILD_COMMANDS)
+	$(POST_BUILD_COMMANDS)

meson.build

@@ -229,6 +229,7 @@ endif
 if get_option('dag_classic')
   files += [
     'src/search/dag_classic/node.cc',
+    'src/search/dag_classic/params.cc',
     'src/search/dag_classic/search.cc',
     'src/search/dag_classic/wrapper.cc',
   ]

proto/net.proto

@@ -397,6 +397,7 @@ message OnnxModel {
   optional string output_wdl = 5;
   optional string output_policy = 6;
   optional string output_mlh = 7;
+  optional string output_err = 8;
 }
 
 message Net {

src/neural/backend.cc

@@ -44,7 +44,7 @@ std::vector<EvalResult> Backend::EvaluateBatch(
     EvalResult& result = results.back();
     result.p.resize(pos.legal_moves.size());
     computation->AddInput(
-        pos, EvalResultPtr{&result.q, &result.d, &result.m,
+        pos, EvalResultPtr{&result.q, &result.d, &result.m, &result.e,
                            std::span<float>(result.p.data(), result.p.size())});
   }
   computation->ComputeBlocking();
@@ -65,4 +65,4 @@ uint64_t Backend::ConfigurationHash(const OptionsDict& options) const {
   return hash;
 }
 
-}  // namespace lczero
+}  // namespace lczero

src/neural/backend.h

@@ -53,17 +53,19 @@ struct EvalResultPtr {
   float* q = nullptr;
   float* d = nullptr;
   float* m = nullptr;
+  float* e = nullptr;
   std::span<float> p = {};
 };
 
 struct EvalResult {
   float q;
   float d;
   float m;
+  float e;
   std::vector<float> p;
 
   EvalResultPtr AsPtr() {
-    return EvalResultPtr{.q = &q, .d = &d, .m = &m, .p = p};
+    return EvalResultPtr{.q = &q, .d = &d, .m = &m, .e = &e, .p = p};
   }
 };
 
@@ -133,4 +135,4 @@ class BackendFactory {
   virtual std::unique_ptr<Backend> Create(const OptionsDict&) = 0;
 };
 
-}  // namespace lczero
+}  // namespace lczero

src/neural/backends/cuda/inputs_outputs.h

@@ -71,7 +71,7 @@ struct CudaGraphExec {
 
 template <typename DataType>
 struct InputsOutputs {
-  InputsOutputs(unsigned maxBatchSize, bool wdl, bool moves_left,
+  InputsOutputs(unsigned maxBatchSize, bool wdl, bool wdl_err, bool moves_left,
                 size_t tensor_mem_size = 0, size_t scratch_size = 0,
                 bool cublasDisableTensorCores = false) {
     ReportCUDAErrors(cudaHostAlloc(
@@ -117,6 +117,17 @@ struct InputsOutputs {
                                               cudaEventDisableTiming));
     ReportCUDAErrors(cudaEventCreateWithFlags(&download_done_event_,
                                               cudaEventDisableTiming));
+
+    if (wdl_err) {
+      ReportCUDAErrors(cudaHostAlloc(
+          &op_value_err_mem_, maxBatchSize * sizeof(op_value_err_mem_[0]),
+          cudaHostAllocMapped));
+      ReportCUDAErrors(
+          cudaMalloc(&op_value_err_mem_gpu_,
+                     maxBatchSize * sizeof(op_value_err_mem_gpu_[0])));
+      ReportCUDAErrors(cudaEventCreateWithFlags(&value_err_done_event_,
+                                                cudaEventDisableTiming));
+    }
     if (moves_left) {
       ReportCUDAErrors(cudaHostAlloc(
           &op_moves_left_mem_, maxBatchSize * sizeof(op_moves_left_mem_[0]),
@@ -172,6 +183,11 @@ struct InputsOutputs {
     ReportCUDAErrors(cudaEventDestroy(value_done_event_));
     ReportCUDAErrors(cudaEventDestroy(wdl_download_done_event_));
     ReportCUDAErrors(cudaEventDestroy(download_done_event_));
+    if (op_value_err_mem_) {
+      ReportCUDAErrors(cudaFreeHost(op_value_err_mem_));
+      ReportCUDAErrors(cudaFree(op_value_err_mem_gpu_));
+      ReportCUDAErrors(cudaEventDestroy(value_err_done_event_));
+    }
     if (op_moves_left_mem_ != nullptr) {
       ReportCUDAErrors(cudaFreeHost(op_moves_left_mem_));
       ReportCUDAErrors(cudaFree(op_moves_left_mem_gpu_));
@@ -199,6 +215,7 @@ struct InputsOutputs {
   DataType* input_val_mem_;
   DataType* op_policy_mem_;
   DataType* op_value_mem_;
+  DataType* op_value_err_mem_ = nullptr;
   DataType* op_moves_left_mem_ = nullptr;
 
   // Copies in VRAM.
@@ -207,6 +224,7 @@ struct InputsOutputs {
   DataType* op_policy_mem_gpu_;
   DataType* op_value_mem_gpu_;
   DataType* op_moves_left_mem_gpu_ = nullptr;
+  DataType* op_value_err_mem_gpu_ = nullptr;
 
   std::unique_ptr<float[]> wdl_cpu_softmax_;
 
@@ -227,6 +245,7 @@ struct InputsOutputs {
   cudaEvent_t upload_done_event_ = nullptr;
   cudaEvent_t policy_done_event_ = nullptr;
   cudaEvent_t value_done_event_ = nullptr;
+  cudaEvent_t value_err_done_event_ = nullptr;
   cudaEvent_t moves_left_done_event_ = nullptr;
   cudaEvent_t wdl_download_done_event_ = nullptr;
   cudaEvent_t download_done_event_ = nullptr;

src/neural/backends/cuda/layers.cc

@@ -2319,13 +2319,14 @@ template <typename DataType>
 ValueHead<DataType>::ValueHead(BaseLayer<DataType>* ip,
                                const MultiHeadWeights::ValueHead& weights,
                                void* scratch, bool attention_body, bool wdl,
-                               ActivationFunction act, int /*max_batch_size*/,
-                               bool use_gemm_ex)
+                               bool wdl_err, ActivationFunction act,
+                               int /*max_batch_size*/, bool use_gemm_ex)
     : BaseLayer<DataType>(weights.ip_val_b.size(), 8, 8, ip),
       embedding_size_(attention_body ? weights.ip_val_b.size()
                                      : weights.value.biases.size()),
       value_hidden_size_(weights.ip1_val_b.size()),
       wdl_(wdl),
+      wdl_err_(wdl_err),
       attention_body_(attention_body),
       act_(act) {
   if (attention_body_) {
@@ -2344,6 +2345,11 @@ ValueHead<DataType>::ValueHead(BaseLayer<DataType>* ip,
 
   allocAndUpload<DataType>(&ip2_val_w_, weights.ip2_val_w, scratch);
   allocAndUpload<DataType>(&ip2_val_b_, weights.ip2_val_b, scratch);
+
+  if (wdl_err_) {
+    allocAndUpload<DataType>(&ip_val_err_w_, weights.ip_val_err_w, scratch);
+    allocAndUpload<DataType>(&ip_val_err_b_, weights.ip_val_err_b, scratch);
+  }
 }
 
 template <typename DataType>
@@ -2356,6 +2362,10 @@ ValueHead<DataType>::~ValueHead() {
   ReportCUDAErrors(cudaFree(ip1_val_b_));
   ReportCUDAErrors(cudaFree(ip2_val_w_));
   ReportCUDAErrors(cudaFree(ip2_val_b_));
+  if (wdl_err_) {
+    ReportCUDAErrors(cudaFree(ip_val_err_w_));
+    ReportCUDAErrors(cudaFree(ip_val_err_b_));
+  }
 }
 
 template <typename DataType>
@@ -2397,7 +2407,7 @@ void ValueHead<DataType>::Eval(int N, DataType* output, const DataType* input,
                              num_outputs, act_, stream);
   }
 
-  {
+  if (!wdl_err_) {
     // Value dense 2
     const int num_inputs = value_hidden_size_;
     const int num_outputs = wdl_ ? 3 : 1;
@@ -2410,6 +2420,17 @@ void ValueHead<DataType>::Eval(int N, DataType* output, const DataType* input,
     addVectors(layer_out, layer_out, ip2_val_b_, num_outputs * batch,
                num_outputs * batch, num_outputs,
                wdl_ ? ACTIVATION_NONE : ACTIVATION_TANH, stream);
+  } else {
+    // Value error dense
+    const int num_inputs = value_hidden_size_;
+    const int num_outputs = 1;
+    const int batch = N;
+    cublasXgemm<DataType>(cublas, CUBLAS_OP_T, CUBLAS_OP_N, num_outputs, batch,
+                          num_inputs, 1.0f, (const DataType*)ip_val_err_w_,
+                          num_inputs, (DataType*)scratch, num_inputs, 0.0f,
+                          output, num_outputs);
+    addVectors(output, output, ip_val_err_b_, N, N, 1, ACTIVATION_SIGMOID,
+               stream);
   }
 }
 

src/neural/backends/cuda/layers.h

@@ -531,7 +531,7 @@ class ValueHead : public BaseLayer<DataType> {
 
  public:
   ValueHead(BaseLayer<DataType>* ip, const MultiHeadWeights::ValueHead& weights,
-            void* scratch, bool attention_body, bool wdl,
+            void* scratch, bool attention_body, bool wdl, bool wdl_err,
             ActivationFunction act, int max_batch_size, bool use_gemm_ex);
   ~ValueHead();
   void Eval(int N, DataType* output, const DataType* input,
@@ -552,6 +552,7 @@ class ValueHead : public BaseLayer<DataType> {
   int embedding_size_;
   int value_hidden_size_;
   bool wdl_;
+  bool wdl_err_;
   bool attention_body_;
   ActivationFunction act_;
 };

src/neural/backends/cuda/network_cuda.cc

@@ -129,7 +129,7 @@ static size_t getMaxAttentionBodySize(const MultiHeadWeights& weights, int N) {
 template <typename DataType>
 class CudaNetworkComputation : public NetworkComputation {
  public:
-  CudaNetworkComputation(CudaNetwork<DataType>* network, bool wdl,
+  CudaNetworkComputation(CudaNetwork<DataType>* network, bool wdl, bool wdl_err,
                          bool moves_left);
   ~CudaNetworkComputation();
 
@@ -177,6 +177,13 @@ class CudaNetworkComputation : public NetworkComputation {
     return 0.0f;
   }
 
+  float GetEVal(int sample) const override {
+    if (wdl_err_) {
+      return FromType(inputs_outputs_->op_value_err_mem_[sample]);
+    }
+    return 0.0f;
+  }
+
   float GetPVal(int sample, int move_id) const override {
     return FromType(
         inputs_outputs_->op_policy_mem_[sample * kNumOutputPolicy + move_id]);
@@ -194,6 +201,7 @@ class CudaNetworkComputation : public NetworkComputation {
   std::unique_ptr<InputsOutputs<DataType>> inputs_outputs_;
   int batch_size_;
   bool wdl_;
+  bool wdl_err_;
   bool moves_left_;
 
   CudaNetwork<DataType>* network_;
@@ -569,9 +577,21 @@ class CudaNetwork : public Network {
              pblczero::NetworkFormat::VALUE_WDL;
       BaseLayer<DataType>* lastlayer = attn_body_ ? encoder_last_ : resi_last_;
       auto value_main = std::make_unique<ValueHead<DataType>>(
-          lastlayer, head, scratch_mem_, attn_body_, wdl_, act, max_batch_size_,
-          use_gemm_ex);
+          lastlayer, head, scratch_mem_, attn_body_, wdl_, false, act,
+          max_batch_size_, use_gemm_ex);
       network_.emplace_back(std::move(value_main));
+
+      if (weights.value_heads.count("st") != 0) {
+        const MultiHeadWeights::ValueHead& st_head =
+            weights.value_heads.at("st");
+        wdl_err_ = st_head.ip_val_err_b.size() > 0;
+        if (wdl_err_) {
+          auto value_err = std::make_unique<ValueHead<DataType>>(
+              lastlayer, st_head, scratch_mem_, attn_body_, wdl_, true, act,
+              max_batch_size_, use_gemm_ex);
+          network_.emplace_back(std::move(value_err));
+        }
+      }
     }
 
     // Moves left head
@@ -634,7 +654,7 @@ class CudaNetwork : public Network {
     // pre-allocate cuda graphs for search threads
     auto allocateCudaGraphs = [&] {
       ReportCUDAErrors(cudaSetDevice(gpu_id_));
-      CudaNetworkComputation<DataType> comp(this, wdl_, moves_left_);
+      CudaNetworkComputation<DataType> comp(this, wdl_, wdl_err_, moves_left_);
       comp.AddInput(InputPlanes{(size_t)kNumInputPlanes});
       // Make sure cublas is initialized in this thread.
       comp.ComputeBlocking();
@@ -775,6 +795,7 @@ class CudaNetwork : public Network {
 
     auto* opPol = io->op_policy_mem_gpu_;
     auto* opVal = io->op_value_mem_gpu_;
+    auto* opValErr = io->op_value_err_mem_gpu_;
     auto* opMov = io->op_moves_left_mem_gpu_;
 
     // Figure out if the memory requirment for running the res block would fit
@@ -919,7 +940,7 @@ class CudaNetwork : public Network {
     network_[l++]->Eval(batchSize, (DataType*)opVal, flow, spare2, scratch_mem,
                         scratch_size_, nullptr, cublas,
                         compute_stream);  // value head
-    if (!moves_left_ && !multi_stream_) {
+    if (!(moves_left_ || wdl_err_) && !multi_stream_) {
 #if CUDA_GRAPH_SUPPORTS_EXTERNAL_EVENTS
       ReportCUDAErrors(
           cudaEventRecordWithFlags(compute_ordering_event_, compute_stream,
@@ -942,6 +963,29 @@ class CudaNetwork : public Network {
 #endif
     }
 
+    if (wdl_err_) {
+      // value error head
+      network_[l++]->Eval(batchSize, (DataType*)opValErr, flow, spare2,
+                          scratch_mem, scratch_size_, nullptr, cublas,
+                          compute_stream);  // value error head
+
+      if (!moves_left_ && !multi_stream_) {
+#if CUDA_GRAPH_SUPPORTS_EXTERNAL_EVENTS
+        ReportCUDAErrors(
+            cudaEventRecordWithFlags(compute_ordering_event_, compute_stream,
+                                     capture ? cudaEventRecordExternal : 0));
+#endif
+      }
+      ReportCUDAErrors(
+          cudaEventRecord(io->value_err_done_event_, compute_stream));
+      ReportCUDAErrors(
+          cudaStreamWaitEvent(download_stream, io->value_err_done_event_, 0));
+      ReportCUDAErrors(
+          cudaMemcpyAsync(io->op_value_err_mem_, io->op_value_err_mem_gpu_,
+                          sizeof(io->op_value_err_mem_[0]) * batchSize,
+                          cudaMemcpyDeviceToHost, download_stream));
+    }
+
     if (moves_left_) {
       // Moves left head
       network_[l++]->Eval(batchSize, spare1, flow, nullptr, scratch_mem,
@@ -1059,15 +1103,16 @@ class CudaNetwork : public Network {
     if (device != gpu_id_) {
       ReportCUDAErrors(cudaSetDevice(gpu_id_));
     }
-    return std::make_unique<CudaNetworkComputation<DataType>>(this, wdl_,
-                                                              moves_left_);
+    return std::make_unique<CudaNetworkComputation<DataType>>(
+        this, wdl_, wdl_err_, moves_left_);
   }
 
   std::unique_ptr<InputsOutputs<DataType>> GetInputsOutputs() {
     std::lock_guard<std::mutex> lock(inputs_outputs_lock_);
     if (free_inputs_outputs_.empty()) {
       return std::make_unique<InputsOutputs<DataType>>(
-          max_batch_size_, wdl_, moves_left_, tensor_mem_size_, scratch_size_,
+          max_batch_size_, wdl_, wdl_err_, moves_left_, tensor_mem_size_,
+          scratch_size_,
           !has_tensor_cores_ && std::is_same<half, DataType>::value);
     } else {
       std::unique_ptr<InputsOutputs<DataType>> resource =
@@ -1098,6 +1143,7 @@ class CudaNetwork : public Network {
   int min_batch_size_;
   bool enable_graph_capture_;
   bool wdl_;
+  bool wdl_err_ = false;
   bool moves_left_;
   bool use_res_block_winograd_fuse_opt_;  // fuse operations inside the residual
                                           // tower
@@ -1213,8 +1259,8 @@ class CudaNetwork : public Network {
 
 template <typename DataType>
 CudaNetworkComputation<DataType>::CudaNetworkComputation(
-    CudaNetwork<DataType>* network, bool wdl, bool moves_left)
-    : wdl_(wdl), moves_left_(moves_left), network_(network) {
+    CudaNetwork<DataType>* network, bool wdl, bool wdl_err, bool moves_left)
+    : wdl_(wdl), wdl_err_(wdl_err), moves_left_(moves_left), network_(network) {
   batch_size_ = 0;
   inputs_outputs_ = network_->GetInputsOutputs();
 }

src/neural/backends/cuda/network_cudnn.cc

@@ -992,7 +992,7 @@ class CudnnNetwork : public Network {
   std::unique_ptr<InputsOutputs<DataType>> GetInputsOutputs() {
     std::lock_guard<std::mutex> lock(inputs_outputs_lock_);
     if (free_inputs_outputs_.empty()) {
-      return std::make_unique<InputsOutputs<DataType>>(max_batch_size_, wdl_,
+      return std::make_unique<InputsOutputs<DataType>>(max_batch_size_, wdl_, false,
                                                        moves_left_);
     } else {
       std::unique_ptr<InputsOutputs<DataType>> resource =