diff --git a/services/surfaceflinger/RenderEngine/Description.cpp b/services/surfaceflinger/RenderEngine/Description.cpp
index c218e4d..09414fd 100644
--- a/services/surfaceflinger/RenderEngine/Description.cpp
+++ b/services/surfaceflinger/RenderEngine/Description.cpp
@@ -51,6 +51,10 @@
     mProjectionMatrix = mtx;
 }
 
+void Description::setSaturationMatrix(const mat4& mtx) {
+    mSaturationMatrix = mtx;
+}
+
 void Description::setColorMatrix(const mat4& mtx) {
     mColorMatrix = mtx;
 }
@@ -78,6 +82,11 @@
     return mColorMatrix != identity;
 }
 
+bool Description::hasSaturationMatrix() const {
+    const mat4 identity;
+    return mSaturationMatrix != identity;
+}
+
 const mat4& Description::getColorMatrix() const {
     return mColorMatrix;
 }
diff --git a/services/surfaceflinger/RenderEngine/Description.h b/services/surfaceflinger/RenderEngine/Description.h
index 6ebb340..06eaf35 100644
--- a/services/surfaceflinger/RenderEngine/Description.h
+++ b/services/surfaceflinger/RenderEngine/Description.h
@@ -42,12 +42,14 @@
     void disableTexture();
     void setColor(const half4& color);
     void setProjectionMatrix(const mat4& mtx);
+    void setSaturationMatrix(const mat4& mtx);
     void setColorMatrix(const mat4& mtx);
     void setInputTransformMatrix(const mat3& matrix);
     void setOutputTransformMatrix(const mat4& matrix);
     bool hasInputTransformMatrix() const;
     bool hasOutputTransformMatrix() const;
     bool hasColorMatrix() const;
+    bool hasSaturationMatrix() const;
     const mat4& getColorMatrix() const;
 
     void setY410BT2020(bool enable);
@@ -90,6 +92,7 @@
     // projection matrix
     mat4 mProjectionMatrix;
     mat4 mColorMatrix;
+    mat4 mSaturationMatrix;
     mat3 mInputTransformMatrix;
     mat4 mOutputTransformMatrix;
 };
diff --git a/services/surfaceflinger/RenderEngine/GLES20RenderEngine.cpp b/services/surfaceflinger/RenderEngine/GLES20RenderEngine.cpp
index 64095dd..aca6c7b 100644
--- a/services/surfaceflinger/RenderEngine/GLES20RenderEngine.cpp
+++ b/services/surfaceflinger/RenderEngine/GLES20RenderEngine.cpp
@@ -263,6 +263,10 @@
     mState.setColorMatrix(colorTransform);
 }
 
+void GLES20RenderEngine::setSaturationMatrix(const mat4& saturationMatrix) {
+    mState.setSaturationMatrix(saturationMatrix);
+}
+
 void GLES20RenderEngine::disableTexturing() {
     mState.disableTexture();
 }
@@ -372,10 +376,11 @@
 
         // we need to convert the RGB value to linear space and convert it back when:
         // - there is a color matrix that is not an identity matrix, or
+        // - there is a saturation matrix that is not an identity matrix, or
         // - there is an output transform matrix that is not an identity matrix, or
         // - the input transfer function doesn't match the output transfer function.
-        if (wideColorState.hasColorMatrix() || wideColorState.hasOutputTransformMatrix() ||
-            inputTransfer != outputTransfer) {
+        if (wideColorState.hasColorMatrix() || wideColorState.hasSaturationMatrix() ||
+            wideColorState.hasOutputTransformMatrix() || inputTransfer != outputTransfer) {
             switch (inputTransfer) {
                 case Dataspace::TRANSFER_ST2084:
                     wideColorState.setInputTransferFunction(Description::TransferFunction::ST2084);
diff --git a/services/surfaceflinger/RenderEngine/GLES20RenderEngine.h b/services/surfaceflinger/RenderEngine/GLES20RenderEngine.h
index c9e402d..84a4813 100644
--- a/services/surfaceflinger/RenderEngine/GLES20RenderEngine.h
+++ b/services/surfaceflinger/RenderEngine/GLES20RenderEngine.h
@@ -81,6 +81,7 @@
     virtual void setupLayerBlackedOut();
     virtual void setupFillWithColor(float r, float g, float b, float a);
     virtual void setupColorTransform(const mat4& colorTransform);
+    virtual void setSaturationMatrix(const mat4& saturationMatrix);
     virtual void disableTexturing();
     virtual void disableBlending();
 
diff --git a/services/surfaceflinger/RenderEngine/Program.cpp b/services/surfaceflinger/RenderEngine/Program.cpp
index fd2c968..95adaca 100644
--- a/services/surfaceflinger/RenderEngine/Program.cpp
+++ b/services/surfaceflinger/RenderEngine/Program.cpp
@@ -135,13 +135,22 @@
         glUniform4fv(mColorLoc, 1, color);
     }
     if (mInputTransformMatrixLoc >= 0) {
-        glUniformMatrix3fv(mInputTransformMatrixLoc, 1, GL_FALSE,
-                           desc.mInputTransformMatrix.asArray());
+        // If the input transform matrix is not identity matrix, we want to merge
+        // the saturation matrix with input transform matrix so that the saturation
+        // matrix is applied at the correct stage.
+        mat4 inputTransformMatrix = mat4(desc.mInputTransformMatrix) * desc.mSaturationMatrix;
+        glUniformMatrix4fv(mInputTransformMatrixLoc, 1, GL_FALSE, inputTransformMatrix.asArray());
     }
     if (mOutputTransformMatrixLoc >= 0) {
         // The output transform matrix and color matrix can be combined as one matrix
         // that is applied right before applying OETF.
         mat4 outputTransformMatrix = desc.mColorMatrix * desc.mOutputTransformMatrix;
+        // If there is no input transform matrix, we want to merge the saturation
+        // matrix with output transform matrix to avoid extra matrix multiplication
+        // in shader.
+        if (mInputTransformMatrixLoc < 0) {
+            outputTransformMatrix *= desc.mSaturationMatrix;
+        }
         glUniformMatrix4fv(mOutputTransformMatrixLoc, 1, GL_FALSE,
                            outputTransformMatrix.asArray());
     }
diff --git a/services/surfaceflinger/RenderEngine/ProgramCache.cpp b/services/surfaceflinger/RenderEngine/ProgramCache.cpp
index abb0290..7c2cb5b 100644
--- a/services/surfaceflinger/RenderEngine/ProgramCache.cpp
+++ b/services/surfaceflinger/RenderEngine/ProgramCache.cpp
@@ -129,7 +129,8 @@
                  description.hasInputTransformMatrix() ?
                      Key::INPUT_TRANSFORM_MATRIX_ON : Key::INPUT_TRANSFORM_MATRIX_OFF)
             .set(Key::Key::OUTPUT_TRANSFORM_MATRIX_MASK,
-                 description.hasOutputTransformMatrix() || description.hasColorMatrix() ?
+                 description.hasOutputTransformMatrix() || description.hasColorMatrix() ||
+                 (!description.hasInputTransformMatrix() && description.hasSaturationMatrix()) ?
                      Key::OUTPUT_TRANSFORM_MATRIX_ON : Key::OUTPUT_TRANSFORM_MATRIX_OFF);
 
     needs.set(Key::Y410_BT2020_MASK,
@@ -309,20 +310,20 @@
                                 float y2 = y1 + (maxOutLumi - y1) * 0.75;
 
                                 // horizontal distances between the last three control points
-                                const float h12 = x2 - x1;
-                                const float h23 = maxInLumi - x2;
+                                float h12 = x2 - x1;
+                                float h23 = maxInLumi - x2;
                                 // tangents at the last three control points
-                                const float m1 = (y2 - y1) / h12;
-                                const float m3 = (maxOutLumi - y2) / h23;
-                                const float m2 = (m1 + m3) / 2.0;
+                                float m1 = (y2 - y1) / h12;
+                                float m3 = (maxOutLumi - y2) / h23;
+                                float m2 = (m1 + m3) / 2.0;
 
                                 if (nits < x0) {
                                     // scale [0.0, x0] to [0.0, y0] linearly
-                                    const float slope = y0 / x0;
+                                    float slope = y0 / x0;
                                     nits *= slope;
                                 } else if (nits < x1) {
                                     // scale [x0, x1] to [y0, y1] linearly
-                                    const float slope = (y1 - y0) / (x1 - x0);
+                                    float slope = (y1 - y0) / (x1 - x0);
                                     nits = y0 + (nits - x0) * slope;
                                 } else if (nits < x2) {
                                     // scale [x1, x2] to [y1, y2] using Hermite interp
@@ -518,10 +519,10 @@
         }
 
         if (needs.hasInputTransformMatrix()) {
-            fs << "uniform mat3 inputTransformMatrix;";
+            fs << "uniform mat4 inputTransformMatrix;";
             fs << R"__SHADER__(
                 highp vec3 InputTransform(const highp vec3 color) {
-                    return inputTransformMatrix * color;
+                    return vec3(inputTransformMatrix * vec4(color, 1.0));
                 }
             )__SHADER__";
         } else {
diff --git a/services/surfaceflinger/RenderEngine/RenderEngine.h b/services/surfaceflinger/RenderEngine/RenderEngine.h
index d559464..a14acaa 100644
--- a/services/surfaceflinger/RenderEngine/RenderEngine.h
+++ b/services/surfaceflinger/RenderEngine/RenderEngine.h
@@ -113,6 +113,7 @@
     virtual void setupFillWithColor(float r, float g, float b, float a) = 0;
 
     virtual void setupColorTransform(const mat4& /* colorTransform */) = 0;
+    virtual void setSaturationMatrix(const mat4& /* saturationMatrix */) = 0;
 
     virtual void disableTexturing() = 0;
     virtual void disableBlending() = 0;
@@ -225,6 +226,7 @@
     void checkErrors() const override;
 
     void setupColorTransform(const mat4& /* colorTransform */) override {}
+    void setSaturationMatrix(const mat4& /* saturationMatrix */) override {}
 
     // internal to RenderEngine
     EGLDisplay getEGLDisplay() const;
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 588d24c..c687e9d 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -2963,10 +2963,8 @@
     ATRACE_INT("hasClientComposition", hasClientComposition);
 
     bool applyColorMatrix = false;
-    bool applyLegacyColorMatrix = false;
-    mat4 colorMatrix;
-    mat4 legacyColorMatrix;
-    const mat4* currentColorMatrix = nullptr;
+    bool needsLegacyColorMatrix = false;
+    bool legacyColorMatrixApplied = false;
 
     if (hasClientComposition) {
         ALOGV("hasClientComposition");
@@ -2985,19 +2983,12 @@
 
         applyColorMatrix = !hasDeviceComposition && !skipClientColorTransform;
         if (applyColorMatrix) {
-            colorMatrix = mDrawingState.colorMatrix;
+            getRenderEngine().setupColorTransform(mDrawingState.colorMatrix);
         }
 
-        applyLegacyColorMatrix = (mDisplayColorSetting == DisplayColorSetting::ENHANCED &&
+        needsLegacyColorMatrix = (mDisplayColorSetting == DisplayColorSetting::ENHANCED &&
                 outputDataspace != Dataspace::UNKNOWN &&
                 outputDataspace != Dataspace::SRGB);
-        if (applyLegacyColorMatrix) {
-            if (applyColorMatrix) {
-                legacyColorMatrix = colorMatrix * mLegacySrgbSaturationMatrix;
-            } else {
-                legacyColorMatrix = mLegacySrgbSaturationMatrix;
-            }
-        }
 
         if (!displayDevice->makeCurrent()) {
             ALOGW("DisplayDevice::makeCurrent failed. Aborting surface composition for display %s",
@@ -3085,18 +3076,14 @@
                 }
                 case HWC2::Composition::Client: {
                     // switch color matrices lazily
-                    if (layer->isLegacyDataSpace()) {
-                        if (applyLegacyColorMatrix && currentColorMatrix != &legacyColorMatrix) {
-                            // TODO(b/78891890) Legacy sRGB saturation matrix should be set
-                            // separately.
-                            getRenderEngine().setupColorTransform(legacyColorMatrix);
-                            currentColorMatrix = &legacyColorMatrix;
+                    if (layer->isLegacyDataSpace() && needsLegacyColorMatrix) {
+                        if (!legacyColorMatrixApplied) {
+                            getRenderEngine().setSaturationMatrix(mLegacySrgbSaturationMatrix);
+                            legacyColorMatrixApplied = true;
                         }
-                    } else {
-                        if (applyColorMatrix && currentColorMatrix != &colorMatrix) {
-                            getRenderEngine().setupColorTransform(colorMatrix);
-                            currentColorMatrix = &colorMatrix;
-                        }
+                    } else if (legacyColorMatrixApplied) {
+                        getRenderEngine().setSaturationMatrix(mat4());
+                        legacyColorMatrixApplied = false;
                     }
 
                     layer->draw(renderArea, clip);
@@ -3111,9 +3098,12 @@
         firstLayer = false;
     }
 
-    if (applyColorMatrix || applyLegacyColorMatrix) {
+    if (applyColorMatrix) {
         getRenderEngine().setupColorTransform(mat4());
     }
+    if (needsLegacyColorMatrix && legacyColorMatrixApplied) {
+        getRenderEngine().setSaturationMatrix(mat4());
+    }
 
     // disable scissor at the end of the frame
     getBE().mRenderEngine->disableScissor();
diff --git a/services/surfaceflinger/tests/unittests/mock/RenderEngine/MockRenderEngine.h b/services/surfaceflinger/tests/unittests/mock/RenderEngine/MockRenderEngine.h
index 93769a5..ac08293 100644
--- a/services/surfaceflinger/tests/unittests/mock/RenderEngine/MockRenderEngine.h
+++ b/services/surfaceflinger/tests/unittests/mock/RenderEngine/MockRenderEngine.h
@@ -61,6 +61,7 @@
     MOCK_METHOD0(setupLayerBlackedOut, void());
     MOCK_METHOD4(setupFillWithColor, void(float, float, float, float));
     MOCK_METHOD1(setupColorTransform, void(const mat4&));
+    MOCK_METHOD1(setSaturationMatrix, void(const mat4&));
     MOCK_METHOD0(disableTexturing, void());
     MOCK_METHOD0(disableBlending, void());
     MOCK_METHOD1(setSourceY410BT2020, void(bool));