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();