diff --git a/include/linux/msm_ion_of.h b/include/linux/msm_ion_of.h
index 1492eca8451f..d1252ff621a9 100644
--- a/include/linux/msm_ion_of.h
+++ b/include/linux/msm_ion_of.h
@@ -3,9 +3,22 @@
  * Copyright (c) 2019, The Linux Foundation. All rights reserved.
  */
 
-#ifndef _MSM_ION_KERNEL_H
-#define _MSM_ION_KERNEL_H
+#ifndef _MSM_ION_OF_H
+#define _MSM_ION_OF_H
+
+#include <linux/device.h>
+#include <uapi/linux/msm_ion.h>
+
+#if IS_ENABLED(CONFIG_ION_MSM_HEAPS)
 
 struct device *msm_ion_heap_device_by_id(int heap_id);
 
-#endif /* _MSM_ION_H */
+#else
+
+static inline struct device *msm_ion_heap_device_by_id(int heap_id)
+{
+	return ERR_PTR(-ENODEV);
+}
+
+#endif /* CONFIG_ION_MSM_HEAPS */
+#endif /* _MSM_ION_OF_H */