//
 // tegra210_amx.c - Tegra210 AMX driver
 //
-// Copyright (c) 2021 NVIDIA CORPORATION.  All rights reserved.
+// Copyright (c) 2021-2023 NVIDIA CORPORATION.  All rights reserved.
 
 #include <linux/clk.h>
 #include <linux/device.h>
        else
                enabled = amx->byte_mask[0] & (1 << reg);
 
+       /*
+        * TODO: Simplify this logic to just return from bytes_map[]
+        *
+        * Presently below is required since bytes_map[] is
+        * tightly packed and cannot store the control value of 256.
+        * Byte mask state is used to know if 256 needs to be returned.
+        * Note that for control value of 256, the put() call stores 0
+        * in the bytes_map[] and disables the corresponding bit in
+        * byte_mask[].
+        */
        if (enabled)
                ucontrol->value.integer.value[0] = bytes_map[reg];
        else
-               ucontrol->value.integer.value[0] = 0;
+               ucontrol->value.integer.value[0] = 256;
 
        return 0;
 }
        unsigned char *bytes_map = (unsigned char *)&amx->map;
        int reg = mc->reg;
        int value = ucontrol->value.integer.value[0];
+       unsigned int mask_val = amx->byte_mask[reg / 32];
 
-       if (value == bytes_map[reg])
+       if (value >= 0 && value <= 255)
+               mask_val |= (1 << (reg % 32));
+       else
+               mask_val &= ~(1 << (reg % 32));
+
+       if (mask_val == amx->byte_mask[reg / 32])
                return 0;
 
-       if (value >= 0 && value <= 255) {
-               /* Update byte map and enable slot */
-               bytes_map[reg] = value;
-               if (reg > 31)
-                       amx->byte_mask[1] |= (1 << (reg - 32));
-               else
-                       amx->byte_mask[0] |= (1 << reg);
-       } else {
-               /* Reset byte map and disable slot */
-               bytes_map[reg] = 0;
-               if (reg > 31)
-                       amx->byte_mask[1] &= ~(1 << (reg - 32));
-               else
-                       amx->byte_mask[0] &= ~(1 << reg);
-       }
+       /* Update byte map and slot */
+       bytes_map[reg] = value % 256;
+       amx->byte_mask[reg / 32] = mask_val;
 
        return 1;
 }