From 088db9f5a8164525f00685c72f00a4baed97b90a Mon Sep 17 00:00:00 2001
From: gfxVPLsdm <gfxvplsdm@intel.com>
Date: Mon, 11 Mar 2024 15:21:25 +0800
Subject: [PATCH] refine (#6535)

Co-authored-by: Zhu, Peak <peak.zhu@intel.com>
---
 .../agnostic/base/hevcehw_base_legacy.cpp     | 58 ++++++++++++++-----
 1 file changed, 44 insertions(+), 14 deletions(-)

diff --git a/_studio/mfx_lib/encode_hw/hevc/agnostic/base/hevcehw_base_legacy.cpp b/_studio/mfx_lib/encode_hw/hevc/agnostic/base/hevcehw_base_legacy.cpp
index 9294328b07..e2647e9646 100644
--- a/_studio/mfx_lib/encode_hw/hevc/agnostic/base/hevcehw_base_legacy.cpp
+++ b/_studio/mfx_lib/encode_hw/hevc/agnostic/base/hevcehw_base_legacy.cpp
@@ -4299,22 +4299,52 @@ mfxStatus Legacy::CheckTU(const ENCODE_CAPS_HEVC& caps, mfxU16& tu)
     if (!tu)
         return MFX_ERR_NONE;
 
-    auto support = caps.TUSupport;
-    mfxI16 abs_diff = 0;
-    bool   sign = 0;
-    mfxI16 newtu = tu;
-
-    do
+    const uint8_t HEVC_EXTENDED_TU_SUPPORT = (1 << (1 - 1)) | (1 << (2 - 1)) | (1 << (4 - 1)) | (1 << (6 - 1)) | (1 << (7 - 1));
+    if (HEVC_EXTENDED_TU_SUPPORT == caps.TUSupport)
     {
-        newtu = tu + (1 - 2 * sign) * abs_diff;
-        abs_diff += !sign;
-        sign = !sign;
-    } while (newtu > 0 && !(support & (1 << (newtu - 1))));
-
-    if (tu != newtu)
+        switch (tu)
+        {
+        case 1:
+            tu = 1;
+            break;
+        case 2:
+            tu = 2;
+            break;
+        case 3: 
+        case 4: 
+        case 5:
+            tu = 4;
+            break;
+        case 6:
+            tu = 6;
+            break;
+        case 7:
+            tu = 7;
+            break;
+        default:
+            tu = 4;
+            break;
+        }
+    }
+    else
     {
-        tu = newtu;
-        return MFX_WRN_INCOMPATIBLE_VIDEO_PARAM;
+        auto support = caps.TUSupport;
+        mfxI16 abs_diff = 0;
+        bool   sign = 0;
+        mfxI16 newtu = tu;
+
+        do
+        {
+            newtu = tu + (1 - 2 * sign) * abs_diff;
+            abs_diff += !sign;
+            sign = !sign;
+        } while (newtu > 0 && !(support & (1 << (newtu - 1))));
+
+        if (tu != newtu)
+        {
+            tu = newtu;
+            return MFX_WRN_INCOMPATIBLE_VIDEO_PARAM;
+        }
     }
 
     return MFX_ERR_NONE;