diff --git a/tf-psa-crypto/drivers/builtin/src/md.c b/tf-psa-crypto/drivers/builtin/src/md.c
index db78f2d..97e11c4 100644
--- a/tf-psa-crypto/drivers/builtin/src/md.c
+++ b/tf-psa-crypto/drivers/builtin/src/md.c
@@ -124,7 +124,7 @@
 };
 #endif
 
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
 static const mbedtls_md_info_t mbedtls_sha3_384_info = {
     MD_INFO(MBEDTLS_MD_SHA3_384, 48, 104)
 };
@@ -175,7 +175,7 @@
         case MBEDTLS_MD_SHA3_256:
             return &mbedtls_sha3_256_info;
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
         case MBEDTLS_MD_SHA3_384:
             return &mbedtls_sha3_384_info;
 #endif
@@ -816,7 +816,7 @@
     MBEDTLS_MD_SHA3_256,
 #endif
 
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
     MBEDTLS_MD_SHA3_384,
 #endif
 
@@ -866,7 +866,7 @@
 #if defined(PSA_WANT_ALG_SHA3_256)
     { "SHA3-256", MBEDTLS_MD_SHA3_256 },
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
     { "SHA3-384", MBEDTLS_MD_SHA3_384 },
 #endif
 #if defined(PSA_WANT_ALG_SHA3_512)
diff --git a/tf-psa-crypto/drivers/builtin/src/oid.c b/tf-psa-crypto/drivers/builtin/src/oid.c
index cf065f1..a182066 100644
--- a/tf-psa-crypto/drivers/builtin/src/oid.c
+++ b/tf-psa-crypto/drivers/builtin/src/oid.c
@@ -773,7 +773,7 @@
         MBEDTLS_MD_SHA3_256,
     },
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_DIGEST_ALG_SHA3_384,    "id-sha3-384",    "SHA-3-384"),
         MBEDTLS_MD_SHA3_384,
@@ -851,12 +851,12 @@
         MBEDTLS_MD_SHA3_256,
     },
 #endif /* PSA_WANT_ALG_SHA3_256 */
-#if defined(MBEDTLS_MD_CAN_SHA3_384)
+#if defined(PSA_WANT_ALG_SHA3_384)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_HMAC_SHA3_384,    "hmacSHA3-384",    "HMAC-SHA3-384"),
         MBEDTLS_MD_SHA3_384,
     },
-#endif /* MBEDTLS_MD_CAN_SHA3_384 */
+#endif /* PSA_WANT_ALG_SHA3_384 */
 #if defined(PSA_WANT_ALG_SHA3_512)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_HMAC_SHA3_512,    "hmacSHA3-512",    "HMAC-SHA3-512"),