diff --git a/crypto/api.c b/crypto/api.c
index b4728811ce3b..959c4e5f264f 100644
--- a/crypto/api.c
+++ b/crypto/api.c
@@ -66,7 +66,8 @@ static inline struct crypto_alg *crypto_alg_mod_lookup(const char *name)
 
 static int crypto_init_flags(struct crypto_tfm *tfm, u32 flags)
 {
-	tfm->crt_flags = 0;
+	tfm->crt_flags = flags & CRYPTO_TFM_REQ_MASK;
+	flags &= ~CRYPTO_TFM_REQ_MASK;
 	
 	switch (crypto_tfm_alg_type(tfm)) {
 	case CRYPTO_ALG_TYPE_CIPHER:
diff --git a/crypto/cipher.c b/crypto/cipher.c
index 8da644364cb4..3df47f93c9db 100644
--- a/crypto/cipher.c
+++ b/crypto/cipher.c
@@ -377,11 +377,7 @@ static int nocrypt_iv(struct crypto_tfm *tfm,
 int crypto_init_cipher_flags(struct crypto_tfm *tfm, u32 flags)
 {
 	u32 mode = flags & CRYPTO_TFM_MODE_MASK;
-	
 	tfm->crt_cipher.cit_mode = mode ? mode : CRYPTO_TFM_MODE_ECB;
-	if (flags & CRYPTO_TFM_REQ_WEAK_KEY)
-		tfm->crt_flags = CRYPTO_TFM_REQ_WEAK_KEY;
-	
 	return 0;
 }
 
diff --git a/crypto/internal.h b/crypto/internal.h
index 37515beafc8c..37aa652ce5ce 100644
--- a/crypto/internal.h
+++ b/crypto/internal.h
@@ -17,6 +17,7 @@
 #include <linux/interrupt.h>
 #include <linux/init.h>
 #include <linux/kernel.h>
+#include <linux/slab.h>
 #include <asm/kmap_types.h>
 
 extern enum km_type crypto_km_types[];
@@ -38,7 +39,7 @@ static inline void crypto_kunmap(void *vaddr, int out)
 
 static inline void crypto_yield(struct crypto_tfm *tfm)
 {
-	if (!in_atomic())
+	if (tfm->crt_flags & CRYPTO_TFM_REQ_MAY_SLEEP)
 		cond_resched();
 }
 
diff --git a/include/linux/crypto.h b/include/linux/crypto.h
index 5e2bcc636a02..3c89df6e7768 100644
--- a/include/linux/crypto.h
+++ b/include/linux/crypto.h
@@ -45,6 +45,7 @@
 #define CRYPTO_TFM_MODE_CTR		0x00000008
 
 #define CRYPTO_TFM_REQ_WEAK_KEY		0x00000100
+#define CRYPTO_TFM_REQ_MAY_SLEEP	0x00000200
 #define CRYPTO_TFM_RES_WEAK_KEY		0x00100000
 #define CRYPTO_TFM_RES_BAD_KEY_LEN   	0x00200000
 #define CRYPTO_TFM_RES_BAD_KEY_SCHED 	0x00400000