diff --git a/Android.bp b/Android.bp
index 872e636f98..49144a7f9b 100644
--- a/Android.bp
+++ b/Android.bp
@@ -13,7 +13,10 @@ display_headers_out = [
 display_kernel_headers_verbose = "--verbose "
 genrule {
     name: "qti_generate_display_kernel_headers",
-    tools: ["headers_install.sh"],
+    tools: [
+        "headers_install.sh",
+        "unifdef"
+    ],
     tool_files: [
         "display_kernel_headers.py",
     ],
@@ -23,6 +26,7 @@ genrule {
          "--header_arch arm64 " +
          "--gen_dir $(genDir) " +
          "--display_include_uapi $(locations include/uapi/*/**/*.h) " +
+         "--unifdef $(location unifdef) " +
          "--headers_install $(location headers_install.sh)",
     out: display_headers_out,
 }
diff --git a/display_kernel_headers.py b/display_kernel_headers.py
index 07be2ea35d..9a7fa15f2a 100644
--- a/display_kernel_headers.py
+++ b/display_kernel_headers.py
@@ -19,31 +19,33 @@ import re
 import subprocess
 import sys
 
-def run_headers_install(verbose, gen_dir, headers_install, prefix, h):
+def run_headers_install(verbose, gen_dir, headers_install, unifdef, prefix, h):
     if not h.startswith(prefix):
         print('error: expected prefix [%s] on header [%s]' % (prefix, h))
         return False
 
     out_h = os.path.join(gen_dir, h[len(prefix):])
     (out_h_dirname, out_h_basename) = os.path.split(out_h)
-    cmd = ["bash", headers_install, h, out_h]
+    env = os.environ.copy()
+    env["LOC_UNIFDEF"] = unifdef
+    cmd = ["sh", headers_install, h, out_h]
 
     if verbose:
         print('run_headers_install: cmd is %s' % cmd)
 
-    result = subprocess.call(cmd)
+    result = subprocess.call(cmd, env=env)
 
     if result != 0:
         print('error: run_headers_install: cmd %s failed %d' % (cmd, result))
         return False
     return True
 
-def gen_display_headers(verbose, gen_dir, headers_install, display_include_uapi):
+def gen_display_headers(verbose, gen_dir, headers_install, unifdef, display_include_uapi):
     error_count = 0
     for h in display_include_uapi:
         display_uapi_include_prefix = os.path.join(h.split('/include/uapi')[0], 'include', 'uapi') + os.sep
         if not run_headers_install(
-                verbose, gen_dir, headers_install,
+                verbose, gen_dir, headers_install, unifdef,
                 display_uapi_include_prefix, h): error_count += 1
     return error_count
 
@@ -69,6 +71,10 @@ def main():
     parser.add_argument(
             '--headers_install', required=True,
             help='The headers_install tool to process input headers.')
+    parser.add_argument(
+              '--unifdef',
+              required=True,
+              help='The unifdef tool used by headers_install.')
 
     args = parser.parse_args()
 
@@ -77,9 +83,10 @@ def main():
         print('gen_dir [%s]' % args.gen_dir)
         print('display_include_uapi [%s]' % args.display_include_uapi)
         print('headers_install [%s]' % args.headers_install)
+        print('unifdef [%s]' % args.unifdef)
 
     return gen_display_headers(args.verbose, args.gen_dir,
-            args.headers_install, args.display_include_uapi)
+            args.headers_install, args.unifdef, args.display_include_uapi)
 
 if __name__ == '__main__':
     sys.exit(main())