Cycles: code refactor to bake using regular render session and tiles

There should be no user visible change from this, except that tile size
now affects performance. The goal here is to simplify bake denoising in
D3099, letting it reuse more denoising tiles and pass code.

A lot of code is now shared with regular rendering, with the two main
differences being that we read some render result passes from the bake API
when starting to render a tile, and call the bake kernel instead of the
path trace kernel.

With this kind of design where Cycles asks for tiles from the bake API,
it should eventually be easier to reduce memory usage, show tiles as
they are baked, or bake multiple passes at once, though there's still
quite some work needed for that.

Reviewers: #cycles

Subscribers: monio, wmatyjewicz, lukasstockner97, michaelknubben

Differential Revision: https://developer.blender.org/D3108

7 files changed, 164 insertions, 280 deletions

diff --git a/intern/cycles/render/bake.cpp b/intern/cycles/render/bake.cpp
index 35f942b3e9b..6044182a51a 100644
--- a/intern/cycles/render/bake.cpp
+++ b/intern/cycles/render/bake.cpp
@@ -15,6 +15,7 @@
  */
 
 #include "render/bake.h"
+#include "render/buffers.h"
 #include "render/integrator.h"
 #include "render/mesh.h"
 #include "render/object.h"
@@ -24,272 +25,130 @@
 
 CCL_NAMESPACE_BEGIN
 
-BakeData::BakeData(const int object, const size_t tri_offset, const size_t num_pixels)
-    : m_object(object), m_tri_offset(tri_offset), m_num_pixels(num_pixels)
+static int aa_samples(Scene *scene, Object *object, ShaderEvalType type)
 {
-  m_primitive.resize(num_pixels);
-  m_u.resize(num_pixels);
-  m_v.resize(num_pixels);
-  m_dudx.resize(num_pixels);
-  m_dudy.resize(num_pixels);
-  m_dvdx.resize(num_pixels);
-  m_dvdy.resize(num_pixels);
-}
-
-BakeData::~BakeData()
-{
-  m_primitive.clear();
-  m_u.clear();
-  m_v.clear();
-  m_dudx.clear();
-  m_dudy.clear();
-  m_dvdx.clear();
-  m_dvdy.clear();
-}
-
-void BakeData::set(int i, int prim, float uv[2], float dudx, float dudy, float dvdx, float dvdy)
-{
-  m_primitive[i] = (prim == -1 ? -1 : m_tri_offset + prim);
-  m_u[i] = uv[0];
-  m_v[i] = uv[1];
-  m_dudx[i] = dudx;
-  m_dudy[i] = dudy;
-  m_dvdx[i] = dvdx;
-  m_dvdy[i] = dvdy;
-}
-
-void BakeData::set_null(int i)
-{
-  m_primitive[i] = -1;
-}
-
-int BakeData::object()
-{
-  return m_object;
-}
-
-size_t BakeData::size()
-{
-  return m_num_pixels;
-}
+  if (type == SHADER_EVAL_UV || type == SHADER_EVAL_ROUGHNESS) {
+    return 1;
+  }
+  else if (type == SHADER_EVAL_NORMAL) {
+    /* Only antialias normal if mesh has bump mapping. */
+    if (object->geometry) {
+      foreach (Shader *shader, object->geometry->used_shaders) {
+        if (shader->has_bump) {
+          return scene->integrator->aa_samples;
+        }
+      }
+    }
 
-bool BakeData::is_valid(int i)
-{
-  return m_primitive[i] != -1;
+    return 1;
+  }
+  else {
+    return scene->integrator->aa_samples;
+  }
 }
 
-uint4 BakeData::data(int i)
+/* Keep it synced with kernel_bake.h logic */
+static int shader_type_to_pass_filter(ShaderEvalType type, int pass_filter)
 {
-  return make_uint4(m_object, m_primitive[i], __float_as_int(m_u[i]), __float_as_int(m_v[i]));
-}
+  const int component_flags = pass_filter &
+                              (BAKE_FILTER_DIRECT | BAKE_FILTER_INDIRECT | BAKE_FILTER_COLOR);
 
-uint4 BakeData::differentials(int i)
-{
-  return make_uint4(__float_as_int(m_dudx[i]),
-                    __float_as_int(m_dudy[i]),
-                    __float_as_int(m_dvdx[i]),
-                    __float_as_int(m_dvdy[i]));
+  switch (type) {
+    case SHADER_EVAL_AO:
+      return BAKE_FILTER_AO;
+    case SHADER_EVAL_SHADOW:
+      return BAKE_FILTER_DIRECT;
+    case SHADER_EVAL_DIFFUSE:
+      return BAKE_FILTER_DIFFUSE | component_flags;
+    case SHADER_EVAL_GLOSSY:
+      return BAKE_FILTER_GLOSSY | component_flags;
+    case SHADER_EVAL_TRANSMISSION:
+      return BAKE_FILTER_TRANSMISSION | component_flags;
+    case SHADER_EVAL_COMBINED:
+      return pass_filter;
+    default:
+      return 0;
+  }
 }
 
 BakeManager::BakeManager()
 {
-  m_bake_data = NULL;
-  m_is_baking = false;
+  type = SHADER_EVAL_BAKE;
+  pass_filter = 0;
+
   need_update = true;
-  m_shader_limit = 512 * 512;
 }
 
 BakeManager::~BakeManager()
 {
-  if (m_bake_data)
-    delete m_bake_data;
 }
 
 bool BakeManager::get_baking()
 {
-  return m_is_baking;
-}
-
-void BakeManager::set_baking(const bool value)
-{
-  m_is_baking = value;
+  return !object_name.empty();
 }
 
-BakeData *BakeManager::init(const int object, const size_t tri_offset, const size_t num_pixels)
+void BakeManager::set(Scene *scene,
+                      const std::string &object_name_,
+                      ShaderEvalType type_,
+                      int pass_filter_)
 {
-  m_bake_data = new BakeData(object, tri_offset, num_pixels);
-  return m_bake_data;
-}
-
-void BakeManager::set_shader_limit(const size_t x, const size_t y)
-{
-  m_shader_limit = x * y;
-  m_shader_limit = (size_t)pow(2, std::ceil(log(m_shader_limit) / log(2)));
-}
+  object_name = object_name_;
+  type = type_;
+  pass_filter = shader_type_to_pass_filter(type_, pass_filter_);
 
-bool BakeManager::bake(Device *device,
-                       DeviceScene *dscene,
-                       Scene *scene,
-                       Progress &progress,
-                       ShaderEvalType shader_type,
-                       const int pass_filter,
-                       BakeData *bake_data,
-                       float result[])
-{
-  size_t num_pixels = bake_data->size();
-
-  int num_samples = aa_samples(scene, bake_data, shader_type);
+  Pass::add(PASS_BAKE_PRIMITIVE, scene->film->passes);
+  Pass::add(PASS_BAKE_DIFFERENTIAL, scene->film->passes);
 
-  /* calculate the total pixel samples for the progress bar */
-  total_pixel_samples = 0;
-  for (size_t shader_offset = 0; shader_offset < num_pixels; shader_offset += m_shader_limit) {
-    size_t shader_size = (size_t)fminf(num_pixels - shader_offset, m_shader_limit);
-    total_pixel_samples += shader_size * num_samples;
+  if (type == SHADER_EVAL_UV) {
+    /* force UV to be available */
+    Pass::add(PASS_UV, scene->film->passes);
   }
-  progress.reset_sample();
-  progress.set_total_pixel_samples(total_pixel_samples);
-
-  /* needs to be up to date for baking specific AA samples */
-  dscene->data.integrator.aa_samples = num_samples;
-  device->const_copy_to("__data", &dscene->data, sizeof(dscene->data));
-
-  for (size_t shader_offset = 0; shader_offset < num_pixels; shader_offset += m_shader_limit) {
-    size_t shader_size = (size_t)fminf(num_pixels - shader_offset, m_shader_limit);
 
-    /* setup input for device task */
-    device_vector<uint4> d_input(device, "bake_input", MEM_READ_ONLY);
-    uint4 *d_input_data = d_input.alloc(shader_size * 2);
-    size_t d_input_size = 0;
-
-    for (size_t i = shader_offset; i < (shader_offset + shader_size); i++) {
-      d_input_data[d_input_size++] = bake_data->data(i);
-      d_input_data[d_input_size++] = bake_data->differentials(i);
-    }
-
-    if (d_input_size == 0) {
-      m_is_baking = false;
-      return false;
-    }
-
-    /* run device task */
-    device_vector<float4> d_output(device, "bake_output", MEM_READ_WRITE);
-    d_output.alloc(shader_size);
-    d_output.zero_to_device();
-    d_input.copy_to_device();
-
-    DeviceTask task(DeviceTask::SHADER);
-    task.shader_input = d_input.device_pointer;
-    task.shader_output = d_output.device_pointer;
-    task.shader_eval_type = shader_type;
-    task.shader_filter = pass_filter;
-    task.shader_x = 0;
-    task.offset = shader_offset;
-    task.shader_w = d_output.size();
-    task.num_samples = num_samples;
-    task.get_cancel = function_bind(&Progress::get_cancel, &progress);
-    task.update_progress_sample = function_bind(&Progress::add_samples_update, &progress, _1, _2);
-
-    device->task_add(task);
-    device->task_wait();
-
-    if (progress.get_cancel()) {
-      d_input.free();
-      d_output.free();
-      m_is_baking = false;
-      return false;
-    }
-
-    d_output.copy_from_device(0, 1, d_output.size());
-    d_input.free();
-
-    /* read result */
-    int k = 0;
-
-    float4 *offset = d_output.data();
-
-    size_t depth = 4;
-    for (size_t i = shader_offset; i < (shader_offset + shader_size); i++) {
-      size_t index = i * depth;
-      float4 out = offset[k++];
-
-      if (bake_data->is_valid(i)) {
-        for (size_t j = 0; j < 4; j++) {
-          result[index + j] = out[j];
-        }
-      }
-    }
-
-    d_output.free();
+  /* force use_light_pass to be true if we bake more than just colors */
+  if (pass_filter & ~BAKE_FILTER_COLOR) {
+    Pass::add(PASS_LIGHT, scene->film->passes);
   }
 
-  m_is_baking = false;
-  return true;
+  /* create device and update scene */
+  scene->film->tag_update(scene);
+  scene->integrator->tag_update(scene);
+
+  need_update = true;
 }
 
 void BakeManager::device_update(Device * /*device*/,
-                                DeviceScene * /*dscene*/,
-                                Scene * /*scene*/,
-                                Progress &progress)
+                                DeviceScene *dscene,
+                                Scene *scene,
+                                Progress & /* progress */)
 {
   if (!need_update)
     return;
 
-  if (progress.get_cancel())
-    return;
+  KernelIntegrator *kintegrator = &dscene->data.integrator;
+  KernelBake *kbake = &dscene->data.bake;
 
-  need_update = false;
-}
-
-void BakeManager::device_free(Device * /*device*/, DeviceScene * /*dscene*/)
-{
-}
-
-int BakeManager::aa_samples(Scene *scene, BakeData *bake_data, ShaderEvalType type)
-{
-  if (type == SHADER_EVAL_UV || type == SHADER_EVAL_ROUGHNESS) {
-    return 1;
-  }
-  else if (type == SHADER_EVAL_NORMAL) {
-    /* Only antialias normal if mesh has bump mapping. */
-    Object *object = scene->objects[bake_data->object()];
+  kbake->type = type;
+  kbake->pass_filter = pass_filter;
 
-    if (object->geometry) {
-      foreach (Shader *shader, object->geometry->used_shaders) {
-        if (shader->has_bump) {
-          return scene->integrator->aa_samples;
-        }
-      }
+  int object_index = 0;
+  foreach (Object *object, scene->objects) {
+    const Geometry *geom = object->geometry;
+    if (object->name == object_name && geom->type == Geometry::MESH) {
+      kbake->object_index = object_index;
+      kbake->tri_offset = geom->prim_offset;
+      kintegrator->aa_samples = aa_samples(scene, object, type);
+      break;
     }
 
-    return 1;
-  }
-  else {
-    return scene->integrator->aa_samples;
+    object_index++;
   }
+
+  need_update = false;
 }
 
-/* Keep it synced with kernel_bake.h logic */
-int BakeManager::shader_type_to_pass_filter(ShaderEvalType type, const int pass_filter)
+void BakeManager::device_free(Device * /*device*/, DeviceScene * /*dscene*/)
 {
-  const int component_flags = pass_filter &
-                              (BAKE_FILTER_DIRECT | BAKE_FILTER_INDIRECT | BAKE_FILTER_COLOR);
-
-  switch (type) {
-    case SHADER_EVAL_AO:
-      return BAKE_FILTER_AO;
-    case SHADER_EVAL_SHADOW:
-      return BAKE_FILTER_DIRECT;
-    case SHADER_EVAL_DIFFUSE:
-      return BAKE_FILTER_DIFFUSE | component_flags;
-    case SHADER_EVAL_GLOSSY:
-      return BAKE_FILTER_GLOSSY | component_flags;
-    case SHADER_EVAL_TRANSMISSION:
-      return BAKE_FILTER_TRANSMISSION | component_flags;
-    case SHADER_EVAL_COMBINED:
-      return pass_filter;
-    default:
-      return 0;
-  }
 }
 
 CCL_NAMESPACE_END
diff --git a/intern/cycles/render/bake.h b/intern/cycles/render/bake.h
index 88537623efb..93e664c2ab1 100644
--- a/intern/cycles/render/bake.h
+++ b/intern/cycles/render/bake.h
@@ -25,67 +25,23 @@
 
 CCL_NAMESPACE_BEGIN
 
-class BakeData {
- public:
-  BakeData(const int object, const size_t tri_offset, const size_t num_pixels);
-  ~BakeData();
-
-  void set(int i, int prim, float uv[2], float dudx, float dudy, float dvdx, float dvdy);
-  void set_null(int i);
-  int object();
-  size_t size();
-  uint4 data(int i);
-  uint4 differentials(int i);
-  bool is_valid(int i);
-
- private:
-  int m_object;
-  size_t m_tri_offset;
-  size_t m_num_pixels;
-  vector<int> m_primitive;
-  vector<float> m_u;
-  vector<float> m_v;
-  vector<float> m_dudx;
-  vector<float> m_dudy;
-  vector<float> m_dvdx;
-  vector<float> m_dvdy;
-};
-
 class BakeManager {
  public:
   BakeManager();
   ~BakeManager();
 
+  void set(Scene *scene, const std::string &object_name, ShaderEvalType type, int pass_filter);
   bool get_baking();
-  void set_baking(const bool value);
-
-  BakeData *init(const int object, const size_t tri_offset, const size_t num_pixels);
-
-  void set_shader_limit(const size_t x, const size_t y);
-
-  bool bake(Device *device,
-            DeviceScene *dscene,
-            Scene *scene,
-            Progress &progress,
-            ShaderEvalType shader_type,
-            const int pass_filter,
-            BakeData *bake_data,
-            float result[]);
 
   void device_update(Device *device, DeviceScene *dscene, Scene *scene, Progress &progress);
   void device_free(Device *device, DeviceScene *dscene);
 
-  static int shader_type_to_pass_filter(ShaderEvalType type, const int pass_filter);
-  static int aa_samples(Scene *scene, BakeData *bake_data, ShaderEvalType type);
-
   bool need_update;
 
-  size_t total_pixel_samples;
-
  private:
-  BakeData *m_bake_data;
-  bool m_is_baking;
-  size_t m_shader_limit;
+  ShaderEvalType type;
+  int pass_filter;
+  std::string object_name;
 };
 
 CCL_NAMESPACE_END
diff --git a/intern/cycles/render/buffers.cpp b/intern/cycles/render/buffers.cpp
index 2d89fb9ffba..b26366af852 100644
--- a/intern/cycles/render/buffers.cpp
+++ b/intern/cycles/render/buffers.cpp
@@ -459,6 +459,40 @@ bool RenderBuffers::get_pass_rect(
   return false;
 }
 
+bool RenderBuffers::set_pass_rect(PassType type, int components, float *pixels)
+{
+  if (buffer.data() == NULL) {
+    return false;
+  }
+
+  int pass_offset = 0;
+
+  for (size_t j = 0; j < params.passes.size(); j++) {
+    Pass &pass = params.passes[j];
+
+    if (pass.type != type) {
+      pass_offset += pass.components;
+      continue;
+    }
+
+    float *out = buffer.data() + pass_offset;
+    int pass_stride = params.get_passes_size();
+    int size = params.width * params.height;
+
+    assert(pass.components == components);
+
+    for (int i = 0; i < size; i++, out += pass_stride, pixels += components) {
+      for (int j = 0; j < components; j++) {
+        out[j] = pixels[j];
+      }
+    }
+
+    return true;
+  }
+
+  return false;
+}
+
 /* Display Buffer */
 
 DisplayBuffer::DisplayBuffer(Device *device, bool linear)
diff --git a/intern/cycles/render/buffers.h b/intern/cycles/render/buffers.h
index 42efb031843..975bae2239c 100644
--- a/intern/cycles/render/buffers.h
+++ b/intern/cycles/render/buffers.h
@@ -92,6 +92,7 @@ class RenderBuffers {
       const string &name, float exposure, int sample, int components, float *pixels);
   bool get_denoising_pass_rect(
       int offset, float exposure, int sample, int components, float *pixels);
+  bool set_pass_rect(PassType type, int components, float *pixels);
 };
 
 /* Display Buffer
@@ -130,7 +131,7 @@ class DisplayBuffer {
 
 class RenderTile {
  public:
-  typedef enum { PATH_TRACE = (1 << 0), DENOISE = (1 << 1) } Task;
+  typedef enum { PATH_TRACE = (1 << 0), BAKE = (1 << 1), DENOISE = (1 << 2) } Task;
 
   Task task;
   int x, y, w, h;
diff --git a/intern/cycles/render/film.cpp b/intern/cycles/render/film.cpp
index 26eda93fadd..d7cbf4a3581 100644
--- a/intern/cycles/render/film.cpp
+++ b/intern/cycles/render/film.cpp
@@ -196,6 +196,10 @@ void Pass::add(PassType type, vector<Pass> &passes, const char *name)
     case PASS_AOV_VALUE:
       pass.components = 1;
       break;
+    case PASS_BAKE_PRIMITIVE:
+    case PASS_BAKE_DIFFERENTIAL:
+      pass.components = 4;
+      break;
     default:
       assert(false);
       break;
@@ -386,11 +390,13 @@ void Film::device_update(Device *device, DeviceScene *dscene, Scene *scene)
     if (pass.type <= PASS_CATEGORY_MAIN_END) {
       kfilm->pass_flag |= pass_flag;
     }
-    else {
-      assert(pass.type <= PASS_CATEGORY_LIGHT_END);
+    else if (pass.type <= PASS_CATEGORY_LIGHT_END) {
       kfilm->use_light_pass = 1;
       kfilm->light_pass_flag |= pass_flag;
     }
+    else {
+      assert(pass.type <= PASS_CATEGORY_BAKE_END);
+    }
 
     switch (pass.type) {
       case PASS_COMBINED:
@@ -471,6 +477,13 @@ void Film::device_update(Device *device, DeviceScene *dscene, Scene *scene)
         kfilm->pass_volume_direct = kfilm->pass_stride;
         break;
 
+      case PASS_BAKE_PRIMITIVE:
+        kfilm->pass_bake_primitive = kfilm->pass_stride;
+        break;
+      case PASS_BAKE_DIFFERENTIAL:
+        kfilm->pass_bake_differential = kfilm->pass_stride;
+        break;
+
 #ifdef WITH_CYCLES_DEBUG
       case PASS_BVH_TRAVERSED_NODES:
         kfilm->pass_bvh_traversed_nodes = kfilm->pass_stride;
diff --git a/intern/cycles/render/session.cpp b/intern/cycles/render/session.cpp
index f7df81a0601..361a1465aac 100644
--- a/intern/cycles/render/session.cpp
+++ b/intern/cycles/render/session.cpp
@@ -410,7 +410,16 @@ bool Session::acquire_tile(RenderTile &rtile, Device *tile_device, uint tile_typ
   rtile.num_samples = tile_manager.state.num_samples;
   rtile.resolution = tile_manager.state.resolution_divider;
   rtile.tile_index = tile->index;
-  rtile.task = tile->state == Tile::DENOISE ? RenderTile::DENOISE : RenderTile::PATH_TRACE;
+
+  if (tile->state == Tile::DENOISE) {
+    rtile.task = RenderTile::DENOISE;
+  }
+  else if (read_bake_tile_cb) {
+    rtile.task = RenderTile::BAKE;
+  }
+  else {
+    rtile.task = RenderTile::PATH_TRACE;
+  }
 
   tile_lock.unlock();
 
@@ -451,11 +460,20 @@ bool Session::acquire_tile(RenderTile &rtile, Device *tile_device, uint tile_typ
   rtile.buffers = tile->buffers;
   rtile.sample = tile_manager.state.sample;
 
-  /* this will tag tile as IN PROGRESS in blender-side render pipeline,
-   * which is needed to highlight currently rendering tile before first
-   * sample was processed for it
-   */
-  update_tile_sample(rtile);
+  if (read_bake_tile_cb) {
+    /* This will read any passes needed as input for baking. */
+    {
+      thread_scoped_lock tile_lock(tile_mutex);
+      read_bake_tile_cb(rtile);
+    }
+    rtile.buffers->buffer.copy_to_device();
+  }
+  else {
+    /* This will tag tile as IN PROGRESS in blender-side render pipeline,
+     * which is needed to highlight currently rendering tile before first
+     * sample was processed for it. */
+    update_tile_sample(rtile);
+  }
 
   return true;
 }
@@ -484,6 +502,7 @@ void Session::release_tile(RenderTile &rtile, const bool need_denoise)
   bool delete_tile;
 
   if (tile_manager.finish_tile(rtile.tile_index, need_denoise, delete_tile)) {
+    /* Finished tile pixels write. */
     if (write_render_tile_cb && params.progressive_refine == false) {
       write_render_tile_cb(rtile);
     }
@@ -494,6 +513,7 @@ void Session::release_tile(RenderTile &rtile, const bool need_denoise)
     }
   }
   else {
+    /* In progress tile pixels update. */
     if (update_render_tile_cb && params.progressive_refine == false) {
       update_render_tile_cb(rtile, false);
     }
diff --git a/intern/cycles/render/session.h b/intern/cycles/render/session.h
index f06952e8020..2707eed5531 100644
--- a/intern/cycles/render/session.h
+++ b/intern/cycles/render/session.h
@@ -148,6 +148,7 @@ class Session {
 
   function<void(RenderTile &)> write_render_tile_cb;
   function<void(RenderTile &, bool)> update_render_tile_cb;
+  function<void(RenderTile &)> read_bake_tile_cb;
 
   explicit Session(const SessionParams &params);
   ~Session();