diff --git a/secure_fw/spm/spm_api.h b/secure_fw/spm/spm_api.h
index 106b29f..a15434d 100644
--- a/secure_fw/spm/spm_api.h
+++ b/secure_fw/spm/spm_api.h
@@ -192,6 +192,17 @@
 #endif
 
 /**
+ * \brief Get the id of the partition for its index from the db
+ *
+ * \param[in] partition_idx     Partition index
+ *
+ * \return Partition ID for that partition
+ *
+ * \note This function doesn't check if partition_idx is valid.
+ */
+uint32_t tfm_spm_partition_get_partition_id(uint32_t partition_idx);
+
+/**
  * \brief Get the flags associated with a partition
  *
  * \param[in] partition_idx     Partition index
@@ -202,6 +213,7 @@
  */
 uint32_t tfm_spm_partition_get_flags(uint32_t partition_idx);
 
+#ifndef TFM_PSA_API
 /**
  * \brief Get the current runtime data of a partition
  *
@@ -235,17 +247,6 @@
         uint32_t stack_ptr, uint32_t lr);
 
 /**
- * \brief Get the id of the partition for its index from the db
- *
- * \param[in] partition_idx     Partition index
- *
- * \return Partition ID for that partition
- *
- * \note This function doesn't check if partition_idx is valid.
- */
-uint32_t tfm_spm_partition_get_partition_id(uint32_t partition_idx);
-
-/**
  * \brief Set the current state of a partition
  *
  * \param[in] partition_idx  Partition index
@@ -314,13 +315,6 @@
                                            const int32_t *args);
 
 /**
- * \brief Initialize partition database
- *
- * \return Error code \ref spm_err_t
- */
-enum spm_err_t tfm_spm_db_init(void);
-
-/**
  * \brief Execute partition init function
  *
  * \return Error code \ref spm_err_t
@@ -335,6 +329,14 @@
  * \note This function doesn't check if partition_idx is valid.
  */
 void tfm_spm_partition_cleanup_context(uint32_t partition_idx);
+#endif /* !defined(TFM_PSA_API) */
+
+/**
+ * \brief Initialize partition database
+ *
+ * \return Error code \ref spm_err_t
+ */
+enum spm_err_t tfm_spm_db_init(void);
 
 /**
  * \brief Change the privilege mode for partition thread mode.