diff --git a/tf-psa-crypto/drivers/builtin/src/oid.c b/tf-psa-crypto/drivers/builtin/src/oid.c
index 972a859..1408843 100644
--- a/tf-psa-crypto/drivers/builtin/src/oid.c
+++ b/tf-psa-crypto/drivers/builtin/src/oid.c
@@ -633,18 +633,18 @@
 
 static const oid_ecp_grp_algid_t oid_ecp_grp_algid[] =
 {
-#if defined(MBEDTLS_ECP_HAVE_CURVE25519)
+#if defined(PSA_WANT_ECC_MONTGOMERY_255)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_X25519,               "X25519",       "X25519"),
         MBEDTLS_ECP_DP_CURVE25519,
     },
-#endif /* MBEDTLS_ECP_HAVE_CURVE25519 */
-#if defined(MBEDTLS_ECP_HAVE_CURVE448)
+#endif /* PSA_WANT_ECC_MONTGOMERY_255 */
+#if defined(PSA_WANT_ECC_MONTGOMERY_448)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_X448,                 "X448",         "X448"),
         MBEDTLS_ECP_DP_CURVE448,
     },
-#endif /* MBEDTLS_ECP_HAVE_CURVE448 */
+#endif /* PSA_WANT_ECC_MONTGOMERY_448 */
     {
         NULL_OID_DESCRIPTOR,
         MBEDTLS_ECP_DP_NONE,
diff --git a/tf-psa-crypto/drivers/builtin/src/pk_internal.h b/tf-psa-crypto/drivers/builtin/src/pk_internal.h
index e86a3a0..2476b6c 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk_internal.h
+++ b/tf-psa-crypto/drivers/builtin/src/pk_internal.h
@@ -114,9 +114,9 @@
 }
 
 /* Helper for Montgomery curves */
-#if defined(MBEDTLS_ECP_HAVE_CURVE25519) || defined(MBEDTLS_ECP_HAVE_CURVE448)
+#if defined(PSA_WANT_ECC_MONTGOMERY_255) || defined(PSA_WANT_ECC_MONTGOMERY_448)
 #define MBEDTLS_PK_HAVE_RFC8410_CURVES
-#endif /* MBEDTLS_ECP_HAVE_CURVE25519 || MBEDTLS_ECP_DP_CURVE448 */
+#endif /* PSA_WANT_ECC_MONTGOMERY_255 || MBEDTLS_ECP_DP_CURVE448 */
 
 #define MBEDTLS_PK_IS_RFC8410_GROUP_ID(id)  \
     ((id == MBEDTLS_ECP_DP_CURVE25519) || (id == MBEDTLS_ECP_DP_CURVE448))