diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 84b10df..93c9ce4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1348,6 +1348,18 @@
     }
     slot->type = attributes->type;
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    /* Find a slot number. Don't yet mark it as allocated in case
+     * the key creation fails or there is a power failure. */
+    if( *p_drv != NULL )
+    {
+        status = psa_find_se_slot_for_key( attributes, *p_drv,
+                                           &slot->data.se.slot_number );
+        if( status != PSA_SUCCESS )
+            return( status );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
     return( status );
 }
 
@@ -1405,6 +1417,18 @@
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    if( driver != NULL )
+    {
+        status = psa_save_se_persistent_data( driver );
+        if( status != PSA_SUCCESS )
+        {
+            psa_destroy_persistent_key( slot->persistent_storage_id );
+            return( status );
+        }
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
     return( status );
 }
 
diff --git a/library/psa_crypto_se.c b/library/psa_crypto_se.c
index b95b2a5..fb57fc9 100644
--- a/library/psa_crypto_se.c
+++ b/library/psa_crypto_se.c
@@ -130,6 +130,35 @@
     return( PSA_SUCCESS );
 }
 
+psa_status_t psa_find_se_slot_for_key(
+    const psa_key_attributes_t *attributes,
+    psa_se_drv_table_entry_t *driver,
+    psa_key_slot_number_t *slot_number )
+{
+    psa_status_t status;
+    psa_drv_se_allocate_key_t p_allocate = NULL;
+
+    /* If the lifetime is wrong, it's a bug in the library. */
+    if( driver->lifetime != attributes->lifetime )
+        return( PSA_ERROR_CORRUPTION_DETECTED );
+
+    /* If the driver doesn't support key creation in any way, give up now. */
+    if( driver->methods->key_management == NULL )
+        return( PSA_ERROR_NOT_SUPPORTED );
+    p_allocate = driver->methods->key_management->p_allocate;
+
+    /* If the driver doesn't tell us how to allocate a slot, that's
+     * not supported for the time being. */
+    if( p_allocate == NULL )
+        return( PSA_ERROR_NOT_SUPPORTED );
+
+    status = ( *p_allocate )( &driver->context,
+                              driver->internal.persistent_data,
+                              attributes,
+                              slot_number );
+    return( status );
+}
+
 
 
 /****************************************************************/
diff --git a/library/psa_crypto_se.h b/library/psa_crypto_se.h
index a9951e6..02819d9 100644
--- a/library/psa_crypto_se.h
+++ b/library/psa_crypto_se.h
@@ -99,6 +99,21 @@
 psa_drv_se_context_t *psa_get_se_driver_context(
     psa_se_drv_table_entry_t *driver );
 
+/** Find a free slot for a key that is to be created.
+ *
+ * This function calls the relevant method in the driver to find a suitable
+ * slot for a key with the given attributes.
+ *
+ * \param[in] attributes    Metadata about the key that is about to be created.
+ * \param[in] driver        The driver table entry to query.
+ * \param[out] slot_number  On success, a slot number that is free in this
+ *                          secure element.
+ */
+psa_status_t psa_find_se_slot_for_key(
+    const psa_key_attributes_t *attributes,
+    psa_se_drv_table_entry_t *driver,
+    psa_key_slot_number_t *slot_number );
+
 /** Load the persistent data of a secure element driver.
  *
  * \param driver        The driver table entry containing the persistent