diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index a2bd201..10b3f98 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -17,7 +17,6 @@
 package com.android.launcher3;
 
 import static com.android.launcher3.Utilities.dpiFromPx;
-import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_TWO_PANEL_HOME;
 import static com.android.launcher3.util.DisplayController.CHANGE_DENSITY;
 import static com.android.launcher3.util.DisplayController.CHANGE_SUPPORTED_BOUNDS;
@@ -44,6 +43,7 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.model.DeviceGridState;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.DisplayController.Info;
 import com.android.launcher3.util.IntArray;
@@ -66,9 +66,6 @@
     public static final MainThreadInitializedObject<InvariantDeviceProfile> INSTANCE =
             new MainThreadInitializedObject<>(InvariantDeviceProfile::new);
 
-    public static final String KEY_MIGRATION_SRC_WORKSPACE_SIZE = "migration_src_workspace_size";
-    public static final String KEY_MIGRATION_SRC_HOTSEAT_COUNT = "migration_src_hotseat_count";
-
     private static final int DEFAULT_TRUE = -1;
     private static final int DEFAULT_SPLIT_DISPLAY = 2;
 
@@ -188,10 +185,7 @@
         if (!newGridName.equals(gridName)) {
             Utilities.getPrefs(context).edit().putString(KEY_IDP_GRID_NAME, newGridName).apply();
         }
-        Utilities.getPrefs(context).edit()
-                .putInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, numDatabaseHotseatIcons)
-                .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, getPointString(numColumns, numRows))
-                .apply();
+        new DeviceGridState(this).writeToPrefs(context);
 
         DisplayController.INSTANCE.get(context).addChangeListener(
                 (displayContext, info, flags) -> {
diff --git a/src/com/android/launcher3/LauncherBackupAgent.java b/src/com/android/launcher3/LauncherBackupAgent.java
index 140794b..dc533f0 100644
--- a/src/com/android/launcher3/LauncherBackupAgent.java
+++ b/src/com/android/launcher3/LauncherBackupAgent.java
@@ -31,6 +31,6 @@
 
     @Override
     public void onRestoreFinished() {
-        RestoreDbTask.setPending(this, true);
+        RestoreDbTask.setPending(this);
     }
 }
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 95a8d81..49f20c6 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -151,15 +151,7 @@
             mOpenHelper = DatabaseHelper.createDatabaseHelper(
                     getContext(), false /* forMigration */);
 
-            if (RestoreDbTask.isPending(getContext())) {
-                if (!RestoreDbTask.performRestore(getContext(), mOpenHelper,
-                        new BackupManager(getContext()))) {
-                    mOpenHelper.createEmptyDB(mOpenHelper.getWritableDatabase());
-                }
-                // Set is pending to false irrespective of the result, so that it doesn't get
-                // executed again.
-                RestoreDbTask.setPending(getContext(), false);
-            }
+            RestoreDbTask.restoreIfNeeded(getContext(), mOpenHelper);
         }
     }
 
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 7d818d2..b92cf09 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -648,21 +648,6 @@
         handler.sendMessage(msg);
     }
 
-    /**
-     * Parses a string encoded using {@link #getPointString(int, int)}
-     */
-    public static Point parsePoint(String point) {
-        String[] split = point.split(",");
-        return new Point(Integer.parseInt(split[0]), Integer.parseInt(split[1]));
-    }
-
-    /**
-     * Encodes a point to string to that it can be persisted atomically.
-     */
-    public static String getPointString(int x, int y) {
-        return String.format(Locale.ENGLISH, "%d,%d", x, y);
-    }
-
     public static void unregisterReceiverSafely(Context context, BroadcastReceiver receiver) {
         try {
             context.unregisterReceiver(receiver);
diff --git a/src/com/android/launcher3/model/DeviceGridState.java b/src/com/android/launcher3/model/DeviceGridState.java
new file mode 100644
index 0000000..761053d
--- /dev/null
+++ b/src/com/android/launcher3/model/DeviceGridState.java
@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model;
+
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_GRID_SIZE_2;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_GRID_SIZE_3;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_GRID_SIZE_4;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_GRID_SIZE_5;
+
+import android.content.Context;
+import android.content.SharedPreferences;
+import android.text.TextUtils;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.Utilities;
+import com.android.launcher3.logging.StatsLogManager.LauncherEvent;
+
+import java.util.Locale;
+import java.util.Objects;
+
+/**
+ * Utility class representing persisted grid properties.
+ */
+public class DeviceGridState {
+
+    public static final String KEY_WORKSPACE_SIZE = "migration_src_workspace_size";
+    public static final String KEY_HOTSEAT_COUNT = "migration_src_hotseat_count";
+    public static final String KEY_DEVICE_TYPE = "migration_src_device_type";
+
+    public static final int TYPE_PHONE = 0;
+    public static final int TYPE_MULTI_DISPLAY = 1;
+    public static final int TYPE_TABLET = 2;
+
+    private final String mGridSizeString;
+    private final int mNumHotseat;
+    private final int mDeviceType;
+
+    public DeviceGridState(InvariantDeviceProfile idp) {
+        mGridSizeString = String.format(Locale.ENGLISH, "%d,%d", idp.numColumns, idp.numRows);
+        mNumHotseat = idp.numDatabaseHotseatIcons;
+        mDeviceType = idp.supportedProfiles.size() > 2
+                ? TYPE_MULTI_DISPLAY
+                : idp.supportedProfiles.stream().allMatch(dp -> dp.isTablet)
+                        ? TYPE_TABLET
+                        : TYPE_PHONE;
+    }
+
+    public DeviceGridState(Context context) {
+        SharedPreferences prefs = Utilities.getPrefs(context);
+        mGridSizeString = prefs.getString(KEY_WORKSPACE_SIZE, "");
+        mNumHotseat = prefs.getInt(KEY_HOTSEAT_COUNT, -1);
+        mDeviceType = prefs.getInt(KEY_DEVICE_TYPE, TYPE_PHONE);
+    }
+
+    /**
+     * Returns the device type for the grid
+     */
+    public int getDeviceType() {
+        return mDeviceType;
+    }
+
+    /**
+     * Stores the device state to shared preferences
+     */
+    public void writeToPrefs(Context context) {
+        Utilities.getPrefs(context).edit()
+                .putString(KEY_WORKSPACE_SIZE, mGridSizeString)
+                .putInt(KEY_HOTSEAT_COUNT, mNumHotseat)
+                .putInt(KEY_DEVICE_TYPE, mDeviceType)
+                .apply();
+    }
+
+    /**
+     * Returns the logging event corresponding to the grid state
+     */
+    public LauncherEvent getWorkspaceSizeEvent() {
+        if (!TextUtils.isEmpty(mGridSizeString)) {
+            switch (mGridSizeString.charAt(0)) {
+                case '5':
+                    return LAUNCHER_GRID_SIZE_5;
+                case '4':
+                    return LAUNCHER_GRID_SIZE_4;
+                case '3':
+                    return LAUNCHER_GRID_SIZE_3;
+                case '2':
+                    return LAUNCHER_GRID_SIZE_2;
+            }
+        }
+        return null;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        DeviceGridState that = (DeviceGridState) o;
+        return mNumHotseat == that.mNumHotseat
+                && mDeviceType == that.mDeviceType
+                && Objects.equals(mGridSizeString, that.mGridSizeString);
+    }
+}
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
index 3935bcf..e64b25c 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -16,9 +16,6 @@
 
 package com.android.launcher3.model;
 
-import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_HOTSEAT_COUNT;
-import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_WORKSPACE_SIZE;
-import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
 import android.content.ComponentName;
@@ -107,11 +104,7 @@
      * Check given a new IDP, if migration is necessary.
      */
     public static boolean needsToMigrate(Context context, InvariantDeviceProfile idp) {
-        SharedPreferences prefs = Utilities.getPrefs(context);
-        String gridSizeString = getPointString(idp.numColumns, idp.numRows);
-
-        return !gridSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, ""))
-                || idp.numDatabaseHotseatIcons != prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, -1);
+        return !new DeviceGridState(idp).equals(new DeviceGridState(context));
     }
 
     /** See {@link #migrateGridIfNeeded(Context, InvariantDeviceProfile)} */
@@ -124,15 +117,15 @@
 
     /**
      * When migrating the grid for preview, we copy the table
-     * {@link LauncherSettings.Favorites.TABLE_NAME} into
-     * {@link LauncherSettings.Favorites.PREVIEW_TABLE_NAME}, run grid size migration from the
+     * {@link LauncherSettings.Favorites#TABLE_NAME} into
+     * {@link LauncherSettings.Favorites#PREVIEW_TABLE_NAME}, run grid size migration from the
      * former to the later, then use the later table for preview.
      *
      * Similarly when doing the actual grid migration, the former grid option's table
-     * {@link LauncherSettings.Favorites.TABLE_NAME} is copied into the new grid option's
-     * {@link LauncherSettings.Favorites.TMP_TABLE}, we then run the grid size migration algorithm
+     * {@link LauncherSettings.Favorites#TABLE_NAME} is copied into the new grid option's
+     * {@link LauncherSettings.Favorites#TMP_TABLE}, we then run the grid size migration algorithm
      * to migrate the later to the former, and load the workspace from the default
-     * {@link LauncherSettings.Favorites.TABLE_NAME}.
+     * {@link LauncherSettings.Favorites#TABLE_NAME}.
      *
      * @return false if the migration failed.
      */
@@ -147,10 +140,7 @@
         }
 
         SharedPreferences prefs = Utilities.getPrefs(context);
-        String gridSizeString = getPointString(idp.numColumns, idp.numRows);
         HashSet<String> validPackages = getValidPackages(context);
-        int srcHotseatCount = prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
-                idp.numDatabaseHotseatIcons);
 
         if (migrateForPreview) {
             if (!LauncherSettings.Settings.call(
@@ -175,11 +165,11 @@
             DbReader srcReader = new DbReader(t.getDb(),
                     migrateForPreview ? LauncherSettings.Favorites.TABLE_NAME
                             : LauncherSettings.Favorites.TMP_TABLE,
-                    context, validPackages, srcHotseatCount);
+                    context, validPackages);
             DbReader destReader = new DbReader(t.getDb(),
                     migrateForPreview ? LauncherSettings.Favorites.PREVIEW_TABLE_NAME
                             : LauncherSettings.Favorites.TABLE_NAME,
-                    context, validPackages, idp.numDatabaseHotseatIcons);
+                    context, validPackages);
 
             Point targetSize = new Point(idp.numColumns, idp.numRows);
             GridSizeMigrationTaskV2 task = new GridSizeMigrationTaskV2(context, t.getDb(),
@@ -202,10 +192,7 @@
 
             if (!migrateForPreview) {
                 // Save current configuration, so that the migration does not run again.
-                prefs.edit()
-                        .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, gridSizeString)
-                        .putInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, idp.numDatabaseHotseatIcons)
-                        .apply();
+                new DeviceGridState(idp).writeToPrefs(context);
             }
         }
     }
@@ -504,7 +491,6 @@
         private final String mTableName;
         private final Context mContext;
         private final HashSet<String> mValidPackages;
-        private final int mHotseatSize;
         private int mLastScreenId = -1;
 
         private final ArrayList<DbEntry> mHotseatEntries = new ArrayList<>();
@@ -513,12 +499,11 @@
                 new ArrayMap<>();
 
         DbReader(SQLiteDatabase db, String tableName, Context context,
-                HashSet<String> validPackages, int hotseatSize) {
+                HashSet<String> validPackages) {
             mDb = db;
             mTableName = tableName;
             mContext = context;
             mValidPackages = validPackages;
-            mHotseatSize = hotseatSize;
         }
 
         protected ArrayList<DbEntry> loadHotseatEntries() {
@@ -543,11 +528,6 @@
                 entry.itemType = c.getInt(indexItemType);
                 entry.screenId = c.getInt(indexScreen);
 
-                if (entry.screenId >= mHotseatSize) {
-                    entriesToRemove.add(entry.id);
-                    continue;
-                }
-
                 try {
                     // calculate weight
                     switch (entry.itemType) {
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 223f4f1..8d02a4a 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.provider;
 
+import static com.android.launcher3.model.DeviceGridState.TYPE_PHONE;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
 import android.app.backup.BackupManager;
@@ -38,6 +39,7 @@
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
+import com.android.launcher3.model.DeviceGridState;
 import com.android.launcher3.model.GridBackupTable;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -57,7 +59,7 @@
 public class RestoreDbTask {
 
     private static final String TAG = "RestoreDbTask";
-    private static final String RESTORE_TASK_PENDING = "restore_task_pending";
+    private static final String RESTORED_DEVICE_TYPE = "restored_task_pending";
 
     private static final String INFO_COLUMN_NAME = "name";
     private static final String INFO_COLUMN_DEFAULT_VALUE = "dflt_value";
@@ -65,13 +67,33 @@
     private static final String APPWIDGET_OLD_IDS = "appwidget_old_ids";
     private static final String APPWIDGET_IDS = "appwidget_ids";
 
-    public static boolean performRestore(Context context, DatabaseHelper helper,
-            BackupManager backupManager) {
+    /**
+     * Tries to restore the backup DB if needed
+     */
+    public static void restoreIfNeeded(Context context, DatabaseHelper helper) {
+        if (!isPending(context)) {
+            return;
+        }
+        if (!performRestore(context, helper)) {
+            helper.createEmptyDB(helper.getWritableDatabase());
+        }
+
+        // Set is pending to false irrespective of the result, so that it doesn't get
+        // executed again.
+        Utilities.getPrefs(context).edit().remove(RESTORED_DEVICE_TYPE).commit();
+    }
+
+    private static boolean performRestore(Context context, DatabaseHelper helper) {
+        if (new DeviceGridState(LauncherAppState.getIDP(context)).getDeviceType()
+                != Utilities.getPrefs(context).getInt(RESTORED_DEVICE_TYPE, TYPE_PHONE)) {
+            // DO not restore if the device types are different
+            return false;
+        }
         SQLiteDatabase db = helper.getWritableDatabase();
         try (SQLiteTransaction t = new SQLiteTransaction(db)) {
             RestoreDbTask task = new RestoreDbTask();
             task.backupWorkspace(context, db);
-            task.sanitizeDB(helper, db, backupManager);
+            task.sanitizeDB(helper, db, new BackupManager(context));
             task.restoreAppWidgetIdsIfExists(context);
             t.commit();
             return true;
@@ -279,12 +301,17 @@
     }
 
     public static boolean isPending(Context context) {
-        return Utilities.getPrefs(context).getBoolean(RESTORE_TASK_PENDING, false);
+        return Utilities.getPrefs(context).contains(RESTORED_DEVICE_TYPE);
     }
 
-    public static void setPending(Context context, boolean isPending) {
-        FileLog.d(TAG, "Restore data received through full backup " + isPending);
-        Utilities.getPrefs(context).edit().putBoolean(RESTORE_TASK_PENDING, isPending).commit();
+    /**
+     * Marks the DB state as pending restoration
+     */
+    public static void setPending(Context context) {
+        FileLog.d(TAG, "Restore data received through full backup ");
+        Utilities.getPrefs(context).edit()
+                .putInt(RESTORED_DEVICE_TYPE, new DeviceGridState(context).getDeviceType())
+                .commit();
     }
 
     private void restoreAppWidgetIdsIfExists(Context context) {