diff --git a/crypto/api.c b/crypto/api.c
index b472881..959c4e5 100644
--- a/crypto/api.c
+++ b/crypto/api.c
@@ -66,7 +66,8 @@
 
 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 8da6443..3df47f9 100644
--- a/crypto/cipher.c
+++ b/crypto/cipher.c
@@ -377,11 +377,7 @@
 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 37515be..37aa652 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_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 5e2bcc6..3c89df6 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