Moves left head for OpenCL backend (LeelaChessZero#1146)

Ttl · web-flow · commit f757103d91f4 · 2020-03-25T06:25:00.000+02:00
diff --git a/src/neural/opencl/OpenCL.h b/src/neural/opencl/OpenCL.h
@@ -68,6 +68,7 @@ class Layer {
   bool is_policy{false};
   bool is_conv_policy{false};
   bool is_value{false};
+  bool is_moves_left{false};
   std::vector<cl::Buffer> weights;
 };
 
@@ -187,6 +188,24 @@ class OpenCL_Network {
     m_layers[layer].ip_out_size = ip_out;
   }
 
+  void push_moves_left(unsigned int channels, unsigned int outputs,
+                       unsigned int ip_in, unsigned int ip_out,
+                       const std::vector<float>& weights,
+                       const std::vector<float>& biases,
+                       const std::vector<float>& fc_w,
+                       const std::vector<float>& fc_b) {
+    size_t layer = get_layer_count();
+    push_weights(layer, weights);
+    push_weights(layer, biases);
+    push_weights(layer, fc_w);
+    push_weights(layer, fc_b);
+    m_layers[layer].is_moves_left = true;
+    m_layers[layer].outputs = outputs;
+    m_layers[layer].channels = channels;
+    m_layers[layer].ip_in_size = ip_in;
+    m_layers[layer].ip_out_size = ip_out;
+  }
+
   size_t get_layer_count() const { return m_layers.size(); }
 
  private:
diff --git a/src/neural/opencl/OpenCLBuffers.cc b/src/neural/opencl/OpenCLBuffers.cc
@@ -46,13 +46,23 @@ OpenCLBuffers::OpenCLBuffers(const OpenCL_Network& opencl_net)
   constexpr auto width = 8;
   constexpr auto height = 8;
 
-  auto finalSize_pol = layers[layers.size() - 2].ip_out_size * sizeof(net_t);
-  auto finalSize_val = layers.back().ip_out_size * sizeof(net_t);
+  m_finalSize_pol = 0;
+  m_finalSize_val = 0;
+  m_finalSize_mov = 0;
 
   auto max_channels = unsigned{0};
   for (const auto& layer : layers) {
     max_channels =
         std::max(max_channels, std::max(layer.channels, layer.outputs));
+    if (layer.is_policy || layer.is_conv_policy) {
+      m_finalSize_pol = layer.ip_out_size * sizeof(net_t);
+    }
+    if (layer.is_value) {
+      m_finalSize_val = layer.ip_out_size * sizeof(net_t);
+    }
+    if (layer.is_moves_left) {
+      m_finalSize_mov = layer.ip_out_size * sizeof(net_t);
+    }
   }
 
   const auto mwg = m_opencl.m_sgemm_tuners.mwg;
@@ -86,16 +96,35 @@ OpenCLBuffers::OpenCLBuffers(const OpenCL_Network& opencl_net)
   try {
     m_pinnedOutBuffer_pol = cl::Buffer(
         m_opencl.m_context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR,
-        max_batch_size * finalSize_pol);
+        max_batch_size * m_finalSize_pol);
   } catch (const cl::Error& e) {
     CERR << "Error in m_pinnedOutBuffer_pol: " << e.what() << ": " << e.err()
          << std::endl;
     throw;
   }
 
-  m_pinnedOutBuffer_val =
-      cl::Buffer(m_opencl.m_context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR,
-                 max_batch_size * finalSize_val);
+  try {
+    m_pinnedOutBuffer_val = cl::Buffer(
+        m_opencl.m_context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR,
+        max_batch_size * m_finalSize_val);
+  } catch (const cl::Error& e) {
+    CERR << "Error in m_pinnedOutBuffer_val: " << e.what() << ": " << e.err()
+         << std::endl;
+    throw;
+  }
+
+  if (m_finalSize_mov > 0) {
+    try {
+      m_pinnedOutBuffer_mov = cl::Buffer(
+          m_opencl.m_context, CL_MEM_WRITE_ONLY | CL_MEM_ALLOC_HOST_PTR,
+          max_batch_size * m_finalSize_mov);
+    } catch (const cl::Error& e) {
+      CERR << "Error in m_pinnedOutBuffer_mov: " << e.what() << ": " << e.err()
+           << std::endl;
+      throw;
+    }
+  }
+
   m_pool_buffer =
       cl::Buffer(m_opencl.m_context, CL_MEM_READ_WRITE | CL_MEM_HOST_NO_ACCESS,
                  alloc_pool_size);
@@ -104,12 +133,10 @@ OpenCLBuffers::OpenCLBuffers(const OpenCL_Network& opencl_net)
 void OpenCLBuffers::forward(const std::vector<net_t>& input,
                             std::vector<net_t>& output_pol,
                             std::vector<net_t>& output_val,
+                            std::vector<net_t>& output_mov,
                             const int batch_size) {
   auto& layers = m_opencl_net.m_layers;
 
-  auto finalSize_pol = layers[layers.size() - 2].ip_out_size * sizeof(net_t);
-  auto finalSize_val = layers.back().ip_out_size * sizeof(net_t);
-
   const auto inSize = sizeof(net_t) * input.size();
   m_commandqueue.enqueueWriteBuffer(m_inBuffer, CL_FALSE, 0, inSize,
                                     input.data());
@@ -241,13 +268,15 @@ void OpenCLBuffers::forward(const std::vector<net_t>& input,
                 layer.outputs * 8 * 8, layer.ip_in_size, layer.ip_out_size);
 
     } else {
-      assert(layer.is_value || layer.is_policy);
+      assert(layer.is_value || layer.is_policy || layer.is_moves_left);
 
       cl::Buffer out_buffer;
       if (layer.is_policy) {
         out_buffer = m_pinnedOutBuffer_pol;
-      } else {
+      } else if (layer.is_value) {
         out_buffer = m_pinnedOutBuffer_val;
+      } else {
+        out_buffer = m_pinnedOutBuffer_mov;
       }
 
       auto conv_weights = begin(layer.weights);
@@ -265,22 +294,36 @@ void OpenCLBuffers::forward(const std::vector<net_t>& input,
 
   auto pinnedOutBufferHost_pol = m_commandqueue.enqueueMapBuffer(
       m_pinnedOutBuffer_pol, CL_FALSE, CL_MAP_READ, 0,
-      batch_size * finalSize_pol);
+      batch_size * m_finalSize_pol);
   auto pinnedOutBufferHost_val = m_commandqueue.enqueueMapBuffer(
       m_pinnedOutBuffer_val, CL_FALSE, CL_MAP_READ, 0,
-      batch_size * finalSize_val);
+      batch_size * m_finalSize_val);
+  void* pinnedOutBufferHost_mov;
+  if (m_finalSize_mov > 0) {
+    pinnedOutBufferHost_mov = m_commandqueue.enqueueMapBuffer(
+        m_pinnedOutBuffer_mov, CL_FALSE, CL_MAP_READ, 0,
+        batch_size * m_finalSize_mov);
+  }
 
   m_commandqueue.finish();
 
   std::memcpy(output_pol.data(), pinnedOutBufferHost_pol,
-              batch_size * finalSize_pol);
+              batch_size * m_finalSize_pol);
   std::memcpy(output_val.data(), pinnedOutBufferHost_val,
-              batch_size * finalSize_val);
+              batch_size * m_finalSize_val);
+  if (m_finalSize_mov > 0) {
+    std::memcpy(output_mov.data(), pinnedOutBufferHost_mov,
+                batch_size * m_finalSize_mov);
+  }
 
   m_commandqueue.enqueueUnmapMemObject(m_pinnedOutBuffer_pol,
                                        pinnedOutBufferHost_pol);
   m_commandqueue.enqueueUnmapMemObject(m_pinnedOutBuffer_val,
                                        pinnedOutBufferHost_val);
+  if (m_finalSize_mov > 0) {
+    m_commandqueue.enqueueUnmapMemObject(m_pinnedOutBuffer_mov,
+                                         pinnedOutBufferHost_mov);
+  }
 }
 
 void OpenCLBuffers::convolve3(int channels, int outputs, cl::Buffer& bufferIn,
diff --git a/src/neural/opencl/OpenCLBuffers.h b/src/neural/opencl/OpenCLBuffers.h
@@ -49,7 +49,8 @@ class OpenCLBuffers {
   OpenCLBuffers(const OpenCL_Network& opencl_net);
 
   void forward(const std::vector<net_t>& input, std::vector<net_t>& output_pol,
-               std::vector<net_t>& output_val, const int batch_size);
+               std::vector<net_t>& output_val, std::vector<net_t>& output_mov,
+               const int batch_size);
 
  private:
   using weight_slice_t = std::vector<cl::Buffer>::const_iterator;
@@ -99,4 +100,8 @@ class OpenCLBuffers {
   cl::Buffer m_pool_buffer;
   cl::Buffer m_pinnedOutBuffer_pol;
   cl::Buffer m_pinnedOutBuffer_val;
+  cl::Buffer m_pinnedOutBuffer_mov;
+  size_t m_finalSize_pol;
+  size_t m_finalSize_val;
+  size_t m_finalSize_mov;
 };
diff --git a/src/neural/opencl/network_opencl.cc b/src/neural/opencl/network_opencl.cc
@@ -45,24 +45,33 @@ class OpenCLNetwork;
 struct OpenCLWeights {
   const std::vector<float> ip2_val_w;
   const std::vector<float> ip2_val_b;
+  const std::vector<float> ip2_mov_w;
+  const std::vector<float> ip2_mov_b;
   const size_t num_output_policies = 1858;
   const size_t num_value_channels;
+  const size_t num_moves_channels;
 
   OpenCLWeights(const WeightsFile& file)
       : ip2_val_w(LayerAdapter(file.weights().ip2_val_w()).as_vector()),
         ip2_val_b(LayerAdapter(file.weights().ip2_val_b()).as_vector()),
-        num_value_channels(LayerAdapter(file.weights().ip1_val_b()).size()) {}
+        ip2_mov_w(LayerAdapter(file.weights().ip2_mov_w()).as_vector()),
+        ip2_mov_b(LayerAdapter(file.weights().ip2_mov_b()).as_vector()),
+        num_value_channels(LayerAdapter(file.weights().ip1_val_b()).size()),
+        num_moves_channels(LayerAdapter(file.weights().ip1_mov_b()).size()) {}
 };
 
 class OpenCLComputation : public NetworkComputation {
  public:
   OpenCLComputation(const OpenCL_Network& opencl_net,
-                    const OpenCLWeights& weights, const bool wdl)
+                    const OpenCLWeights& weights, const bool wdl,
+                    const bool moves_left)
       : opencl_net_(opencl_net),
         weights_(weights),
         policies_(),
         q_values_(),
-        wdl_(wdl) {
+        m_values_(),
+        wdl_(wdl),
+        moves_left_(moves_left) {
     buffers_ = opencl_net.acquire_buffers();
   }
 
@@ -82,6 +91,7 @@ class OpenCLComputation : public NetworkComputation {
 
     const auto num_output_policies = weights_.num_output_policies;
     const auto num_value_channels = weights_.num_value_channels;
+    const auto num_moves_channels = weights_.num_moves_channels;
 
     // Typically
     // input_channels = 112
@@ -90,6 +100,7 @@ class OpenCLComputation : public NetworkComputation {
 
     std::vector<float> output_pol(largest_batch_size * num_output_policies);
     std::vector<float> output_val(largest_batch_size * num_value_channels);
+    std::vector<float> output_mov(largest_batch_size * num_moves_channels);
     std::vector<float> input_data(largest_batch_size * kInputPlanes * kSquares);
 
     for (size_t i = 0; i < plane_count; i += largest_batch_size) {
@@ -98,7 +109,8 @@ class OpenCLComputation : public NetworkComputation {
         EncodePlanes(planes_[i + j], &input_data[j * kSquares * kInputPlanes]);
       }
 
-      buffers_->forward(input_data, output_pol, output_val, batch_size);
+      buffers_->forward(input_data, output_pol, output_val, output_mov,
+                        batch_size);
 
       for (size_t j = 0; j < batch_size; j++) {
         std::vector<float> policy(num_output_policies);
@@ -135,6 +147,16 @@ class OpenCLComputation : public NetworkComputation {
 
           q_values_.emplace_back(std::tanh(winrate));
         }
+
+        if (moves_left_) {
+          auto m = weights_.ip2_mov_b[0];
+          auto ptr_weights = weights_.ip2_mov_w.data();
+          auto ptr_outputs = &output_mov[j * num_moves_channels];
+          for (size_t i = 0; i < num_moves_channels; i++)
+            m += ptr_weights[i] * std::max(0.0f, ptr_outputs[i]);
+
+          m_values_.emplace_back(std::max(0.0f, m));
+        }
       }
     }
   }
@@ -162,8 +184,13 @@ class OpenCLComputation : public NetworkComputation {
     }
   }
 
-  float GetMVal(int /* sample */) const override {
-    return 0.0f;
+  float GetMVal(int sample) const override {
+    if (moves_left_) {
+      auto d = m_values_[sample];
+      return d;
+    } else {
+      return 0.0f;
+    }
   }
 
   // Returns P value @move_id of @sample.
@@ -185,9 +212,11 @@ class OpenCLComputation : public NetworkComputation {
 
   std::vector<std::vector<float>> policies_;
   std::vector<float> q_values_;
+  std::vector<float> m_values_;
 
   std::unique_ptr<OpenCLBuffers> buffers_;
   bool wdl_;
+  bool moves_left_;
 };
 
 void OpenCLComputation::EncodePlanes(const InputPlanes& sample, float* buffer) {
@@ -205,7 +234,7 @@ class OpenCLNetwork : public Network {
 
   OpenCLNetwork(const WeightsFile& file, const OptionsDict& options)
       : capabilities_{file.format().network_format().input(),
-                      pblczero::NetworkFormat::MOVES_LEFT_NONE},
+                      file.format().network_format().moves_left()},
         weights_(file),
         params_(),
         opencl_(),
@@ -222,6 +251,9 @@ class OpenCLNetwork : public Network {
     wdl_ = file.format().network_format().output() ==
            pblczero::NetworkFormat::OUTPUT_WDL;
 
+    moves_left_ = file.format().network_format().moves_left() ==
+                  pblczero::NetworkFormat::MOVES_LEFT_V1;
+
     auto max_batch_size_ =
         static_cast<size_t>(options.GetOrDefault<int>("batch_size", 16));
     if (max_batch_size_ > kHardMaxBatchSize) {
@@ -241,9 +273,11 @@ class OpenCLNetwork : public Network {
     const auto residual_blocks = weights.residual.size();
 
     const auto num_value_input_planes = weights.value.biases.size();
+    const auto num_moves_input_planes = weights.moves_left.biases.size();
     const auto num_policy_input_planes = weights.policy.biases.size();
     const auto num_output_policy = kPolicyOutputs;
     const auto num_value_channels = weights.ip1_val_b.size();
+    const auto num_moves_channels = weights.ip1_mov_b.size();
 
     // Typically
     // input_channels = 112
@@ -350,11 +384,20 @@ class OpenCLNetwork : public Network {
                            weights.value.biases, weights.ip1_val_w,
                            weights.ip1_val_b);
 
+    if (moves_left_) {
+      opencl_net_.push_moves_left(
+          channels, num_moves_input_planes,
+          num_moves_input_planes * width * height, num_moves_channels,
+          weights.moves_left.weights, weights.moves_left.biases,
+          weights.ip1_mov_w, weights.ip1_mov_b);
+    }
+
     opencl_net_.setMaxMatchSize(max_batch_size_);
   }
 
   std::unique_ptr<NetworkComputation> NewComputation() override {
-    return std::make_unique<OpenCLComputation>(opencl_net_, weights_, wdl_);
+    return std::make_unique<OpenCLComputation>(opencl_net_, weights_, wdl_,
+                                               moves_left_);
   }
 
   const NetworkCapabilities& GetCapabilities() const override {
@@ -372,6 +415,7 @@ class OpenCLNetwork : public Network {
   OpenCL opencl_;
   OpenCL_Network opencl_net_;
   bool wdl_;
+  bool moves_left_;
 };
 
 std::unique_ptr<Network> MakeOpenCLNetwork(const WeightsFile& weights,
