diff --git a/include/efi.h b/include/efi.h
index d07187c7dd..3d587807e8 100644
--- a/include/efi.h
+++ b/include/efi.h
@@ -30,8 +30,11 @@ struct efi_device_path;
 
 #define EFI_BITS_PER_LONG	BITS_PER_LONG
 
-/* With 64-bit EFI stub, EFI_BITS_PER_LONG has to be 64 */
-#ifdef CONFIG_EFI_STUB_64BIT
+/*
+ * With 64-bit EFI stub, EFI_BITS_PER_LONG has to be 64. EFI_STUB is set
+ * in lib/efi/Makefile, when building the stub.
+ */
+#if defined(CONFIG_EFI_STUB_64BIT) && defined(EFI_STUB)
 #undef EFI_BITS_PER_LONG
 #define EFI_BITS_PER_LONG	64
 #endif
diff --git a/lib/efi/Makefile b/lib/efi/Makefile
index e32dc14f3d..9449600b5b 100644
--- a/lib/efi/Makefile
+++ b/lib/efi/Makefile
@@ -9,9 +9,9 @@ obj-$(CONFIG_EFI_STUB) += efi_info.o
 
 CFLAGS_REMOVE_efi_stub.o := -mregparm=3 \
 	$(if $(CONFIG_EFI_STUB_64BIT),-march=i386 -m32)
-CFLAGS_efi_stub.o := -fpic -fshort-wchar
+CFLAGS_efi_stub.o := -fpic -fshort-wchar -DEFI_STUB
 CFLAGS_REMOVE_efi.o := -mregparm=3 \
 	$(if $(CONFIG_EFI_STUB_64BIT),-march=i386 -m32)
-CFLAGS_efi.o := -fpic -fshort-wchar
+CFLAGS_efi.o := -fpic -fshort-wchar -DEFI_STUB
 
 extra-$(CONFIG_EFI_STUB) += efi_stub.o efi.o