Multires: Initial Face Sets support

This implements the Sculpt Mode API functions needed for Face Sets and
visibility management for PBVH_GRIDS. No major changes were needed in
the operators and the sculpt mode code. This implementation stores the
face sets in the base mesh, so faces created in higher subdivision
levels can't be modified individually. Also, we are not checking for
multiple face sets per vertex (that can be added in the future), so
relax tools don't work yet. The rest of the features (paint, undo,
visibility operators..) work as expected.

Reviewed By: brecht

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

11 files changed, 196 insertions, 73 deletions

diff --git a/source/blender/blenkernel/BKE_paint.h b/source/blender/blenkernel/BKE_paint.h
index 881f3356a86..8c925ee2ae1 100644
--- a/source/blender/blenkernel/BKE_paint.h
+++ b/source/blender/blenkernel/BKE_paint.h
@@ -296,6 +296,7 @@ typedef struct SculptSession {
   int *pmap_mem;
 
   /* Mesh Face Sets */
+  int totfaces;
   int *face_sets;
 
   /* BMesh for dynamic topology sculpting */
diff --git a/source/blender/blenkernel/BKE_pbvh.h b/source/blender/blenkernel/BKE_pbvh.h
index 16a7e4d38d0..b4f16bfd899 100644
--- a/source/blender/blenkernel/BKE_pbvh.h
+++ b/source/blender/blenkernel/BKE_pbvh.h
@@ -219,6 +219,8 @@ int BKE_pbvh_count_grid_quads(BLI_bitmap **grid_hidden,
                               int totgrid,
                               int gridsize);
 
+void BKE_pbvh_sync_face_sets_to_grids(PBVH *bvh);
+
 /* multires level, only valid for type == PBVH_GRIDS */
 const struct CCGKey *BKE_pbvh_get_grid_key(const PBVH *pbvh);
 
@@ -298,6 +300,8 @@ void BKE_pbvh_grids_update(PBVH *bvh,
                            void **gridfaces,
                            struct DMFlagMat *flagmats,
                            unsigned int **grid_hidden);
+void BKE_pbvh_subdiv_cgg_set(PBVH *bvh, struct SubdivCCG *subdiv_ccg);
+void BKE_pbvh_face_sets_set(PBVH *bvh, int *face_sets);
 
 void BKE_pbvh_face_sets_color_set(PBVH *bvh, int seed, int color_default);
 
diff --git a/source/blender/blenkernel/BKE_subdiv_ccg.h b/source/blender/blenkernel/BKE_subdiv_ccg.h
index f8534371b17..99b134dab3e 100644
--- a/source/blender/blenkernel/BKE_subdiv_ccg.h
+++ b/source/blender/blenkernel/BKE_subdiv_ccg.h
@@ -305,6 +305,8 @@ void BKE_subdiv_ccg_neighbor_coords_get(const SubdivCCG *subdiv_ccg,
                                         const bool include_duplicates,
                                         SubdivCCGNeighbors *r_neighbors);
 
+int BKE_subdiv_cgg_grid_to_face_index(const SubdivCCG *subdiv_ccg, const int grid_index);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/source/blender/blenkernel/intern/paint.c b/source/blender/blenkernel/intern/paint.c
index 719336f7351..f52ec5f568f 100644
--- a/source/blender/blenkernel/intern/paint.c
+++ b/source/blender/blenkernel/intern/paint.c
@@ -1525,6 +1525,7 @@ static void sculpt_update_object(
     ss->multires = mmd;
     ss->totvert = me_eval->totvert;
     ss->totpoly = me_eval->totpoly;
+    ss->totfaces = me->totpoly;
     ss->mvert = NULL;
     ss->mpoly = NULL;
     ss->mloop = NULL;
@@ -1532,25 +1533,26 @@ static void sculpt_update_object(
   else {
     ss->totvert = me->totvert;
     ss->totpoly = me->totpoly;
+    ss->totfaces = me->totpoly;
     ss->mvert = me->mvert;
     ss->mpoly = me->mpoly;
     ss->mloop = me->mloop;
     ss->multires = NULL;
     ss->vmask = CustomData_get_layer(&me->vdata, CD_PAINT_MASK);
+  }
 
-    /* Sculpt Face Sets. */
-    if (!CustomData_has_layer(&me->pdata, CD_SCULPT_FACE_SETS)) {
-      ss->face_sets = CustomData_add_layer(
-          &me->pdata, CD_SCULPT_FACE_SETS, CD_CALLOC, NULL, me->totpoly);
-      for (int i = 0; i < me->totpoly; i++) {
-        ss->face_sets[i] = 1;
-      }
-
-      /* Set the default face set color if the datalayer did not exist. */
-      me->face_sets_color_default = 1;
+  /* Sculpt Face Sets. */
+  if (!CustomData_has_layer(&me->pdata, CD_SCULPT_FACE_SETS)) {
+    ss->face_sets = CustomData_add_layer(
+        &me->pdata, CD_SCULPT_FACE_SETS, CD_CALLOC, NULL, me->totpoly);
+    for (int i = 0; i < me->totpoly; i++) {
+      ss->face_sets[i] = 1;
     }
-    ss->face_sets = CustomData_get_layer(&me->pdata, CD_SCULPT_FACE_SETS);
+
+    /* Set the default face set color if the datalayer did not exist. */
+    me->face_sets_color_default = 1;
   }
+  ss->face_sets = CustomData_get_layer(&me->pdata, CD_SCULPT_FACE_SETS);
 
   ss->subdiv_ccg = me_eval->runtime.subdiv_ccg;
 
@@ -1558,6 +1560,9 @@ static void sculpt_update_object(
   BLI_assert(pbvh == ss->pbvh);
   UNUSED_VARS_NDEBUG(pbvh);
 
+  BKE_pbvh_subdiv_cgg_set(ss->pbvh, ss->subdiv_ccg);
+  BKE_pbvh_face_sets_set(ss->pbvh, ss->face_sets);
+
   BKE_pbvh_face_sets_color_set(ss->pbvh, me->face_sets_color_seed, me->face_sets_color_default);
 
   if (need_pmap && ob->type == OB_MESH && !ss->pmap) {
diff --git a/source/blender/blenkernel/intern/pbvh.c b/source/blender/blenkernel/intern/pbvh.c
index 9a4ce8acb11..ac0c9d030cf 100644
--- a/source/blender/blenkernel/intern/pbvh.c
+++ b/source/blender/blenkernel/intern/pbvh.c
@@ -383,6 +383,25 @@ int BKE_pbvh_count_grid_quads(BLI_bitmap **grid_hidden,
   return totquad;
 }
 
+void BKE_pbvh_sync_face_sets_to_grids(PBVH *bvh)
+{
+  const int gridsize = bvh->gridkey.grid_size;
+  for (int i = 0; i < bvh->totgrid; i++) {
+    BLI_bitmap *gh = bvh->grid_hidden[i];
+    const int face_index = BKE_subdiv_cgg_grid_to_face_index(bvh->subdiv_ccg, i);
+    if (!gh && bvh->face_sets[face_index] < 0) {
+      gh = bvh->grid_hidden[i] = BLI_BITMAP_NEW(bvh->gridkey.grid_area, "partialvis_update_grids");
+    }
+    if (gh) {
+      for (int y = 0; y < gridsize; y++) {
+        for (int x = 0; x < gridsize; x++) {
+          BLI_BITMAP_SET(gh, y * gridsize + x, bvh->face_sets[face_index] < 0);
+        }
+      }
+    }
+  }
+}
+
 static void build_grid_leaf_node(PBVH *bvh, PBVHNode *node)
 {
   int totquads = BKE_pbvh_count_grid_quads(
@@ -1289,10 +1308,14 @@ static void pbvh_update_draw_buffer_cb(void *__restrict userdata,
     switch (bvh->type) {
       case PBVH_GRIDS:
         GPU_pbvh_grid_buffers_update(node->draw_buffers,
+                                     bvh->subdiv_ccg,
                                      bvh->grids,
                                      bvh->grid_flag_mats,
                                      node->prim_indices,
                                      node->totprim,
+                                     bvh->face_sets,
+                                     bvh->face_sets_color_seed,
+                                     bvh->face_sets_color_default,
                                      &bvh->gridkey,
                                      update_flags);
         break;
@@ -2974,3 +2997,13 @@ MVert *BKE_pbvh_get_verts(const PBVH *bvh)
   BLI_assert(bvh->type == PBVH_FACES);
   return bvh->verts;
 }
+
+void BKE_pbvh_subdiv_cgg_set(PBVH *bvh, SubdivCCG *subdiv_ccg)
+{
+  bvh->subdiv_ccg = subdiv_ccg;
+}
+
+void BKE_pbvh_face_sets_set(PBVH *bvh, int *face_sets)
+{
+  bvh->face_sets = face_sets;
+}
diff --git a/source/blender/blenkernel/intern/pbvh_intern.h b/source/blender/blenkernel/intern/pbvh_intern.h
index af92f11e219..21cb5649330 100644
--- a/source/blender/blenkernel/intern/pbvh_intern.h
+++ b/source/blender/blenkernel/intern/pbvh_intern.h
@@ -138,6 +138,7 @@ struct PBVH {
 
   int face_sets_color_seed;
   int face_sets_color_default;
+  int *face_sets;
 
   /* Grid Data */
   CCGKey gridkey;
@@ -168,6 +169,7 @@ struct PBVH {
   int cd_face_node_offset;
 
   struct BMLog *bm_log;
+  struct SubdivCCG *subdiv_ccg;
 };
 
 /* pbvh.c */
diff --git a/source/blender/blenkernel/intern/subdiv_ccg.c b/source/blender/blenkernel/intern/subdiv_ccg.c
index 3c1a9c4d3d6..dff7170e517 100644
--- a/source/blender/blenkernel/intern/subdiv_ccg.c
+++ b/source/blender/blenkernel/intern/subdiv_ccg.c
@@ -1780,3 +1780,14 @@ void BKE_subdiv_ccg_neighbor_coords_get(const SubdivCCG *subdiv_ccg,
   }
 #endif
 }
+
+int BKE_subdiv_cgg_grid_to_face_index(const SubdivCCG *subdiv_ccg, const int grid_index)
+{
+  Subdiv *subdiv = subdiv_ccg->subdiv;
+  OpenSubdiv_TopologyRefiner *topology_refiner = subdiv->topology_refiner;
+  SubdivCCGFace *face = subdiv_ccg->grid_faces[grid_index];
+
+  const int face_grid_index = grid_index - face->start_grid_index;
+  const int face_index = face - subdiv_ccg->faces;
+  return face_index;
+}
diff --git a/source/blender/editors/sculpt_paint/sculpt.c b/source/blender/editors/sculpt_paint/sculpt.c
index 6ccd197c908..50b8c9a9677 100644
--- a/source/blender/editors/sculpt_paint/sculpt.c
+++ b/source/blender/editors/sculpt_paint/sculpt.c
@@ -248,7 +248,8 @@ static void SCULPT_face_set_visibility_set(SculptSession *ss, int face_set, bool
 {
   switch (BKE_pbvh_type(ss->pbvh)) {
     case PBVH_FACES:
-      for (int i = 0; i < ss->totpoly; i++) {
+    case PBVH_GRIDS:
+      for (int i = 0; i < ss->totfaces; i++) {
         if (abs(ss->face_sets[i]) == face_set) {
           if (visible) {
             ss->face_sets[i] = abs(ss->face_sets[i]);
@@ -261,8 +262,6 @@ static void SCULPT_face_set_visibility_set(SculptSession *ss, int face_set, bool
       break;
     case PBVH_BMESH:
       break;
-    case PBVH_GRIDS:
-      break;
   }
 }
 
@@ -270,14 +269,13 @@ static void SCULPT_face_sets_visibility_invert(SculptSession *ss)
 {
   switch (BKE_pbvh_type(ss->pbvh)) {
     case PBVH_FACES:
-      for (int i = 0; i < ss->totpoly; i++) {
+    case PBVH_GRIDS:
+      for (int i = 0; i < ss->totfaces; i++) {
         ss->face_sets[i] *= -1;
       }
       break;
     case PBVH_BMESH:
       break;
-    case PBVH_GRIDS:
-      break;
   }
 }
 
@@ -285,7 +283,8 @@ static void SCULPT_face_sets_visibility_all_set(SculptSession *ss, bool visible)
 {
   switch (BKE_pbvh_type(ss->pbvh)) {
     case PBVH_FACES:
-      for (int i = 0; i < ss->totpoly; i++) {
+    case PBVH_GRIDS:
+      for (int i = 0; i < ss->totfaces; i++) {
 
         /* This can run on geometry without a face set assigned, so its ID sign can't be changed to
          * modify the visibility. Force that geometry to the ID 1 to enable changing the visibility
@@ -304,8 +303,6 @@ static void SCULPT_face_sets_visibility_all_set(SculptSession *ss, bool visible)
       break;
     case PBVH_BMESH:
       break;
-    case PBVH_GRIDS:
-      break;
   }
 }
 
@@ -343,8 +340,12 @@ static bool SCULPT_vertex_all_face_sets_visible_get(SculptSession *ss, int index
     }
     case PBVH_BMESH:
       return true;
-    case PBVH_GRIDS:
-      return true;
+    case PBVH_GRIDS: {
+      const CCGKey *key = BKE_pbvh_get_grid_key(ss->pbvh);
+      const int grid_index = index / key->grid_area;
+      const int face_index = BKE_subdiv_cgg_grid_to_face_index(ss->subdiv_ccg, grid_index);
+      return ss->face_sets[face_index] > 0;
+    }
   }
   return true;
 }
@@ -362,8 +363,15 @@ static void SCULPT_vertex_face_set_set(SculptSession *ss, int index, int face_se
     } break;
     case PBVH_BMESH:
       break;
-    case PBVH_GRIDS:
-      break;
+    case PBVH_GRIDS: {
+      const CCGKey *key = BKE_pbvh_get_grid_key(ss->pbvh);
+      const int grid_index = index / key->grid_area;
+      const int face_index = BKE_subdiv_cgg_grid_to_face_index(ss->subdiv_ccg, grid_index);
+      if (ss->face_sets[face_index] > 0) {
+        ss->face_sets[face_index] = abs(face_set);
+      }
+
+    } break;
   }
 }
 
@@ -382,8 +390,12 @@ int SCULPT_vertex_face_set_get(SculptSession *ss, int index)
     }
     case PBVH_BMESH:
       return 0;
-    case PBVH_GRIDS:
-      return 0;
+    case PBVH_GRIDS: {
+      const CCGKey *key = BKE_pbvh_get_grid_key(ss->pbvh);
+      const int grid_index = index / key->grid_area;
+      const int face_index = BKE_subdiv_cgg_grid_to_face_index(ss->subdiv_ccg, grid_index);
+      return ss->face_sets[face_index];
+    }
   }
   return 0;
 }
@@ -402,8 +414,12 @@ bool SCULPT_vertex_has_face_set(SculptSession *ss, int index, int face_set)
     }
     case PBVH_BMESH:
       return true;
-    case PBVH_GRIDS:
-      return true;
+    case PBVH_GRIDS: {
+      const CCGKey *key = BKE_pbvh_get_grid_key(ss->pbvh);
+      const int grid_index = index / key->grid_area;
+      const int face_index = BKE_subdiv_cgg_grid_to_face_index(ss->subdiv_ccg, grid_index);
+      return ss->face_sets[face_index] == face_set;
+    }
   }
   return true;
 }
@@ -415,8 +431,17 @@ static void sculpt_visibility_sync_face_sets_to_vertex(SculptSession *ss, int in
 
 void SCULPT_visibility_sync_all_face_sets_to_vertices(SculptSession *ss)
 {
-  for (int i = 0; i < ss->totvert; i++) {
-    sculpt_visibility_sync_face_sets_to_vertex(ss, i);
+  switch (BKE_pbvh_type(ss->pbvh)) {
+    case PBVH_FACES: {
+      for (int i = 0; i < ss->totvert; i++) {
+        sculpt_visibility_sync_face_sets_to_vertex(ss, i);
+      }
+    }
+    case PBVH_GRIDS: {
+      BKE_pbvh_sync_face_sets_to_grids(ss->pbvh);
+    }
+    case PBVH_BMESH:
+      break;
   }
 }
 
@@ -439,7 +464,7 @@ static void UNUSED_FUNCTION(sculpt_visibility_sync_vertex_to_face_sets)(SculptSe
 void SCULPT_visibility_sync_all_vertex_to_face_sets(SculptSession *ss)
 {
   if (BKE_pbvh_type(ss->pbvh) == PBVH_FACES) {
-    for (int i = 0; i < ss->totpoly; i++) {
+    for (int i = 0; i < ss->totfaces; i++) {
       MPoly *poly = &ss->mpoly[i];
       bool poly_visible = true;
       for (int l = 0; l < poly->totloop; l++) {
@@ -479,7 +504,7 @@ bool SCULPT_vertex_has_unique_face_set(SculptSession *ss, int index)
     case PBVH_BMESH:
       return false;
     case PBVH_GRIDS:
-      return false;
+      return true;
   }
   return false;
 }
@@ -487,9 +512,10 @@ bool SCULPT_vertex_has_unique_face_set(SculptSession *ss, int index)
 static int SCULPT_face_set_next_available_get(SculptSession *ss)
 {
   switch (BKE_pbvh_type(ss->pbvh)) {
-    case PBVH_FACES: {
+    case PBVH_FACES:
+    case PBVH_GRIDS: {
       int next_face_set = 0;
-      for (int i = 0; i < ss->totpoly; i++) {
+      for (int i = 0; i < ss->totfaces; i++) {
         if (abs(ss->face_sets[i]) > next_face_set) {
           next_face_set = abs(ss->face_sets[i]);
         }
@@ -499,8 +525,6 @@ static int SCULPT_face_set_next_available_get(SculptSession *ss)
     }
     case PBVH_BMESH:
       return 0;
-    case PBVH_GRIDS:
-      return 0;
   }
   return 0;
 }
@@ -10286,7 +10310,7 @@ static void sculpt_mask_expand_cancel(bContext *C, wmOperator *op)
   for (int n = 0; n < ss->filter_cache->totnode; n++) {
     PBVHNode *node = ss->filter_cache->nodes[n];
     if (create_face_set) {
-      for (int i = 0; i < ss->totpoly; i++) {
+      for (int i = 0; i < ss->totfaces; i++) {
         ss->face_sets[i] = ss->filter_cache->prev_face_set[i];
       }
     }
@@ -10492,7 +10516,7 @@ static int sculpt_mask_expand_modal(bContext *C, wmOperator *op, const wmEvent *
   if (mask_expand_update_it < ss->filter_cache->mask_update_last_it) {
 
     if (create_face_set) {
-      for (int i = 0; i < ss->totpoly; i++) {
+      for (int i = 0; i < ss->totfaces; i++) {
         ss->face_sets[i] = ss->filter_cache->prev_face_set[i];
       }
     }
@@ -10619,8 +10643,8 @@ static int sculpt_mask_expand_invoke(bContext *C, wmOperator *op, const wmEvent
   }
 
   if (create_face_set) {
-    ss->filter_cache->prev_face_set = MEM_callocN(sizeof(float) * ss->totpoly, "prev face mask");
-    for (int i = 0; i < ss->totpoly; i++) {
+    ss->filter_cache->prev_face_set = MEM_callocN(sizeof(float) * ss->totfaces, "prev face mask");
+    for (int i = 0; i < ss->totfaces; i++) {
       ss->filter_cache->prev_face_set[i] = ss->face_sets[i];
     }
     ss->filter_cache->new_face_set = SCULPT_face_set_next_available_get(ss);
@@ -11172,8 +11196,8 @@ static int sculpt_face_set_create_invoke(bContext *C, wmOperator *op, const wmEv
 
   const int mode = RNA_enum_get(op->ptr, "mode");
 
-  /* Dyntopo and Multires not supported for now. */
-  if (BKE_pbvh_type(ss->pbvh) != PBVH_FACES) {
+  /* Dyntopo not suported. */
+  if (BKE_pbvh_type(ss->pbvh) == PBVH_BMESH) {
     return OPERATOR_CANCELLED;
   }
 
@@ -11514,8 +11538,8 @@ static int sculpt_face_set_init_invoke(bContext *C, wmOperator *op, const wmEven
 
   const int mode = RNA_enum_get(op->ptr, "mode");
 
-  /* Dyntopo and Multires not supported for now. */
-  if (BKE_pbvh_type(ss->pbvh) != PBVH_FACES) {
+  /* Dyntopo not supported. */
+  if (BKE_pbvh_type(ss->pbvh) == PBVH_BMESH) {
     return OPERATOR_CANCELLED;
   }
 
@@ -11673,8 +11697,8 @@ static int sculpt_face_sets_change_visibility_invoke(bContext *C,
   ARegion *region = CTX_wm_region(C);
   Depsgraph *depsgraph = CTX_data_depsgraph_pointer(C);
 
-  /* Dyntopo and Multires not supported for now. */
-  if (BKE_pbvh_type(ss->pbvh) != PBVH_FACES) {
+  /* Dyntopo not supported. */
+  if (BKE_pbvh_type(ss->pbvh) == PBVH_BMESH) {
     return OPERATOR_CANCELLED;
   }
 
@@ -11714,7 +11738,7 @@ static int sculpt_face_sets_change_visibility_invoke(bContext *C,
       }
     }
 
-    for (int i = 0; i < ss->totpoly; i++) {
+    for (int i = 0; i < ss->totfaces; i++) {
       if (ss->face_sets[i] <= 0) {
         hidden_vertex = true;
         break;
@@ -11834,8 +11858,9 @@ static int sculpt_face_sets_randomize_colors_invoke(bContext *C,
 
   mesh->face_sets_color_seed += 1;
   if (ss->face_sets) {
-    const int random_index = clamp_i(
-        ss->totpoly * BLI_hash_int_01(mesh->face_sets_color_seed), 0, max_ii(0, ss->totpoly - 1));
+    const int random_index = clamp_i(ss->totfaces * BLI_hash_int_01(mesh->face_sets_color_seed),
+                                     0,
+                                     max_ii(0, ss->totfaces - 1));
     mesh->face_sets_color_default = ss->face_sets[random_index];
   }
   BKE_pbvh_face_sets_color_set(pbvh, mesh->face_sets_color_seed, mesh->face_sets_color_default);
diff --git a/source/blender/editors/sculpt_paint/sculpt_undo.c b/source/blender/editors/sculpt_paint/sculpt_undo.c
index af404c64fb0..5233ce8d257 100644
--- a/source/blender/editors/sculpt_paint/sculpt_undo.c
+++ b/source/blender/editors/sculpt_paint/sculpt_undo.c
@@ -48,6 +48,7 @@
 #include "BKE_main.h"
 #include "BKE_mesh.h"
 #include "BKE_multires.h"
+#include "BKE_object.h"
 #include "BKE_paint.h"
 #include "BKE_scene.h"
 #include "BKE_subdiv_ccg.h"
@@ -335,8 +336,11 @@ static bool sculpt_undo_restore_face_sets(bContext *C, SculptUndoNode *unode)
 {
   ViewLayer *view_layer = CTX_data_view_layer(C);
   Object *ob = OBACT(view_layer);
-  SculptSession *ss = ob->sculpt;
-  memcpy(ss->face_sets, unode->face_sets, ss->totpoly * sizeof(int));
+  Mesh *me = BKE_object_get_original_mesh(ob);
+  int *face_sets = CustomData_get_layer(&me->pdata, CD_SCULPT_FACE_SETS);
+  for (int i = 0; i < me->totpoly; i++) {
+    face_sets[i] = unode->face_sets[i];
+  }
   return false;
 }
 
@@ -542,7 +546,6 @@ static void sculpt_undo_restore_list(bContext *C, Depsgraph *depsgraph, ListBase
       return;
     }
     else if (unode->type == SCULPT_UNDO_FACE_SETS) {
-
       sculpt_undo_restore_face_sets(C, unode);
 
       rebuild = true;
@@ -551,14 +554,15 @@ static void sculpt_undo_restore_list(bContext *C, Depsgraph *depsgraph, ListBase
       BKE_sculpt_update_object_for_edit(depsgraph, ob, true, need_mask);
 
       SCULPT_visibility_sync_all_face_sets_to_vertices(ss);
+
       BKE_pbvh_update_vertex_data(ss->pbvh, PBVH_UpdateVisibility);
 
       if (BKE_pbvh_type(ss->pbvh) == PBVH_FACES) {
         BKE_mesh_flush_hidden_from_verts(ob->data);
       }
 
+      DEG_id_tag_update(&ob->id, ID_RECALC_SHADING);
       if (!BKE_sculptsession_use_pbvh_draw(ob, v3d)) {
-        DEG_id_tag_update(&ob->id, ID_RECALC_SHADING);
         DEG_id_tag_update(&ob->id, ID_RECALC_GEOMETRY);
       }
 
@@ -983,8 +987,6 @@ static SculptUndoNode *sculpt_undo_geometry_push(Object *ob, SculptUndoType type
 static SculptUndoNode *sculpt_undo_face_sets_push(Object *ob, SculptUndoType type)
 {
   UndoSculpt *usculpt = sculpt_undo_get_nodes();
-  SculptSession *ss = ob->sculpt;
-
   SculptUndoNode *unode = usculpt->nodes.first;
 
   unode = MEM_callocN(sizeof(*unode), __func__);
@@ -993,8 +995,14 @@ static SculptUndoNode *sculpt_undo_face_sets_push(Object *ob, SculptUndoType typ
   unode->type = type;
   unode->applied = true;
 
-  unode->face_sets = MEM_callocN(ss->totpoly * sizeof(int), "sculpt face sets");
-  memcpy(unode->face_sets, ss->face_sets, ss->totpoly * sizeof(int));
+  Mesh *me = BKE_object_get_original_mesh(ob);
+
+  unode->face_sets = MEM_callocN(me->totpoly * sizeof(int), "sculpt face sets");
+
+  int *face_sets = CustomData_get_layer(&me->pdata, CD_SCULPT_FACE_SETS);
+  for (int i = 0; i < me->totpoly; i++) {
+    unode->face_sets[i] = face_sets[i];
+  }
 
   BLI_addtail(&usculpt->nodes, unode);
 
diff --git a/source/blender/gpu/GPU_buffers.h b/source/blender/gpu/GPU_buffers.h
index 9d91fd79137..b02f6ed09b3 100644
--- a/source/blender/gpu/GPU_buffers.h
+++ b/source/blender/gpu/GPU_buffers.h
@@ -95,10 +95,14 @@ void GPU_pbvh_bmesh_buffers_update(GPU_PBVH_Buffers *buffers,
                                    const int update_flags);
 
 void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
+                                  struct SubdivCCG *subdiv_ccg,
                                   struct CCGElem **grids,
                                   const struct DMFlagMat *grid_flag_mats,
                                   int *grid_indices,
                                   int totgrid,
+                                  const int *sculpt_face_sets,
+                                  const int face_sets_color_seed,
+                                  const int face_sets_color_default,
                                   const struct CCGKey *key,
                                   const int update_flags);
 
diff --git a/source/blender/gpu/intern/gpu_buffers.c b/source/blender/gpu/intern/gpu_buffers.c
index 168a3c83a91..cfc1eb05731 100644
--- a/source/blender/gpu/intern/gpu_buffers.c
+++ b/source/blender/gpu/intern/gpu_buffers.c
@@ -43,6 +43,7 @@
 #include "BKE_mesh.h"
 #include "BKE_paint.h"
 #include "BKE_pbvh.h"
+#include "BKE_subdiv_ccg.h"
 
 #include "GPU_batch.h"
 #include "GPU_buffers.h"
@@ -523,8 +524,13 @@ GPU_PBVH_Buffers *GPU_pbvh_mesh_buffers_build(const int (*face_vert_indices)[3],
 /** \name Grid PBVH
  * \{ */
 
-static void gpu_pbvh_grid_fill_index_buffers(
-    GPU_PBVH_Buffers *buffers, int *grid_indices, uint visible_quad_len, int totgrid, int gridsize)
+static void gpu_pbvh_grid_fill_index_buffers(GPU_PBVH_Buffers *buffers,
+                                             SubdivCCG *subdiv_ccg,
+                                             const int *face_sets,
+                                             int *grid_indices,
+                                             uint visible_quad_len,
+                                             int totgrid,
+                                             int gridsize)
 {
   GPUIndexBufBuilder elb, elb_lines;
   GPUIndexBufBuilder elb_fast, elb_lines_fast;
@@ -594,7 +600,6 @@ static void gpu_pbvh_grid_fill_index_buffers(
     const uint grid_vert_len = square_uint(gridsize - 1) * 4;
     for (int i = 0; i < totgrid; i++, offset += grid_vert_len) {
       bool grid_visible = false;
-
       BLI_bitmap *gh = buffers->grid_hidden[grid_indices[i]];
 
       uint v0, v1, v2, v3;
@@ -673,16 +678,22 @@ void GPU_pbvh_grid_buffers_update_free(GPU_PBVH_Buffers *buffers,
 
 /* Threaded - do not call any functions that use OpenGL calls! */
 void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
+                                  SubdivCCG *subdiv_ccg,
                                   CCGElem **grids,
-                                  const DMFlagMat *grid_flag_mats,
+                                  const struct DMFlagMat *grid_flag_mats,
                                   int *grid_indices,
                                   int totgrid,
-                                  const CCGKey *key,
+                                  const int *sculpt_face_sets,
+                                  const int face_sets_color_seed,
+                                  const int face_sets_color_default,
+                                  const struct CCGKey *key,
                                   const int update_flags)
 {
   const bool show_mask = (update_flags & GPU_PBVH_BUFFERS_SHOW_MASK) != 0;
   const bool show_vcol = (update_flags & GPU_PBVH_BUFFERS_SHOW_VCOL) != 0;
   bool empty_mask = true;
+  bool default_face_set = true;
+
   int i, j, k, x, y;
 
   /* Build VBO */
@@ -702,8 +713,13 @@ void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
       return;
     }
 
-    gpu_pbvh_grid_fill_index_buffers(
-        buffers, grid_indices, visible_quad_len, totgrid, key->grid_size);
+    gpu_pbvh_grid_fill_index_buffers(buffers,
+                                     subdiv_ccg,
+                                     sculpt_face_sets,
+                                     grid_indices,
+                                     visible_quad_len,
+                                     totgrid,
+                                     key->grid_size);
   }
 
   uint vbo_index_offset = 0;
@@ -716,9 +732,23 @@ void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
     }
 
     for (i = 0; i < totgrid; i++) {
-      CCGElem *grid = grids[grid_indices[i]];
+      const int grid_index = grid_indices[i];
+      CCGElem *grid = grids[grid_index];
       int vbo_index = vbo_index_offset;
 
+      uchar face_set_color[4] = {UCHAR_MAX, UCHAR_MAX, UCHAR_MAX, UCHAR_MAX};
+
+      if (subdiv_ccg && sculpt_face_sets) {
+        const int face_index = BKE_subdiv_cgg_grid_to_face_index(subdiv_ccg, grid_index);
+
+        const int fset = abs(sculpt_face_sets[face_index]);
+        /* Skip for the default color Face Set to render it white. */
+        if (fset != face_sets_color_default) {
+          face_set_overlay_color_get(fset, face_sets_color_seed, face_set_color);
+          default_face_set = false;
+        }
+      }
+
       if (buffers->smooth) {
         for (y = 0; y < key->grid_size; y++) {
           for (x = 0; x < key->grid_size; x++) {
@@ -742,8 +772,7 @@ void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
               GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.col, vbo_index, &vcol);
             }
 
-            uchar fsets[3] = {UCHAR_MAX, UCHAR_MAX, UCHAR_MAX};
-            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index, &fsets);
+            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index, &face_set_color);
 
             vbo_index += 1;
           }
@@ -799,11 +828,10 @@ void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
             GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.col, vbo_index + 2, &vcol);
             GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.col, vbo_index + 3, &vcol);
 
-            uchar fsets[3] = {UCHAR_MAX, UCHAR_MAX, UCHAR_MAX};
-            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 0, &fsets);
-            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 1, &fsets);
-            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 2, &fsets);
-            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 3, &fsets);
+            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 0, &face_set_color);
+            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 1, &face_set_color);
+            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 2, &face_set_color);
+            GPU_vertbuf_attr_set(buffers->vert_buf, g_vbo_id.fset, vbo_index + 3, &face_set_color);
 
             vbo_index += 4;
           }
@@ -823,7 +851,7 @@ void GPU_pbvh_grid_buffers_update(GPU_PBVH_Buffers *buffers,
   buffers->totgrid = totgrid;
   buffers->grid_flag_mats = grid_flag_mats;
   buffers->gridkey = *key;
-  buffers->show_overlay = !empty_mask;
+  buffers->show_overlay = !empty_mask || !default_face_set;
 }
 
 /* Threaded - do not call any functions that use OpenGL calls! */