u32 mask, bitmap, ao_bitmap, ao_cu_mask = 0;
        struct amdgpu_cu_info *cu_info = &adev->gfx.cu_info;
        unsigned disable_masks[4 * 2];
+       u32 ao_cu_num;
+
+       if (adev->flags & AMD_IS_APU)
+               ao_cu_num = 2;
+       else
+               ao_cu_num = adev->gfx.config.max_cu_per_sh;
 
        memset(cu_info, 0, sizeof(*cu_info));
 
                        bitmap = gfx_v6_0_get_cu_enabled(adev);
                        cu_info->bitmap[i][j] = bitmap;
 
-                       for (k = 0; k < 16; k++) {
+                       for (k = 0; k < adev->gfx.config.max_cu_per_sh; k++) {
                                if (bitmap & mask) {
-                                       if (counter < 2)
+                                       if (counter < ao_cu_num)
                                                ao_bitmap |= mask;
                                        counter ++;
                                }
 
        u32 mask, bitmap, ao_bitmap, ao_cu_mask = 0;
        struct amdgpu_cu_info *cu_info = &adev->gfx.cu_info;
        unsigned disable_masks[4 * 2];
+       u32 ao_cu_num;
+
+       if (adev->flags & AMD_IS_APU)
+               ao_cu_num = 2;
+       else
+               ao_cu_num = adev->gfx.config.max_cu_per_sh;
 
        memset(cu_info, 0, sizeof(*cu_info));
 
                        bitmap = gfx_v7_0_get_cu_active_bitmap(adev);
                        cu_info->bitmap[i][j] = bitmap;
 
-                       for (k = 0; k < 16; k ++) {
+                       for (k = 0; k < adev->gfx.config.max_cu_per_sh; k ++) {
                                if (bitmap & mask) {
-                                       if (counter < 2)
+                                       if (counter < ao_cu_num)
                                                ao_bitmap |= mask;
                                        counter ++;
                                }
 
        u32 mask, bitmap, ao_bitmap, ao_cu_mask = 0;
        struct amdgpu_cu_info *cu_info = &adev->gfx.cu_info;
        unsigned disable_masks[4 * 2];
+       u32 ao_cu_num;
 
        memset(cu_info, 0, sizeof(*cu_info));
 
+       if (adev->flags & AMD_IS_APU)
+               ao_cu_num = 2;
+       else
+               ao_cu_num = adev->gfx.config.max_cu_per_sh;
+
        amdgpu_gfx_parse_disable_cu(disable_masks, 4, 2);
 
        mutex_lock(&adev->grbm_idx_mutex);
                        bitmap = gfx_v8_0_get_cu_active_bitmap(adev);
                        cu_info->bitmap[i][j] = bitmap;
 
-                       for (k = 0; k < 16; k ++) {
+                       for (k = 0; k < adev->gfx.config.max_cu_per_sh; k ++) {
                                if (bitmap & mask) {
-                                       if (counter < 2)
+                                       if (counter < ao_cu_num)
                                                ao_bitmap |= mask;
                                        counter ++;
                                }
 
                        bitmap = gfx_v9_0_get_cu_active_bitmap(adev);
                        cu_info->bitmap[i][j] = bitmap;
 
-                       for (k = 0; k < 16; k ++) {
+                       for (k = 0; k < adev->gfx.config.max_cu_per_sh; k ++) {
                                if (bitmap & mask) {
-                                       if (counter < 2)
+                                       if (counter < adev->gfx.config.max_cu_per_sh)
                                                ao_bitmap |= mask;
                                        counter ++;
                                }