diff --git a/kernel_headers.py b/kernel_headers.py
index 65ae1e7..27ef474 100644
--- a/kernel_headers.py
+++ b/kernel_headers.py
@@ -309,7 +309,7 @@
   return error_count
 
 
-def run_headers_install(verbose, gen_dir, headers_install, prefix, h):
+def run_headers_install(verbose, gen_dir, headers_install, unifdef, prefix, h):
   """Process a header through the headers_install script.
 
   The headers_install script does some processing of a header so that it is
@@ -324,6 +324,7 @@
     verbose: Set True to print progress messages.
     gen_dir: Where to place the generated files.
     headers_install: The script that munges the header.
+    unifdef: The unifdef tool used by headers_install.
     prefix: The prefix to strip from h to generate the output filename.
     h: The input header to process.
   Return:
@@ -343,7 +344,9 @@
   if verbose:
     print('run_headers_install: cmd is %s' % cmd)
 
-  result = subprocess.call(['sh', headers_install, h, out_h])
+  env = os.environ.copy()
+  env["LOC_UNIFDEF"] = unifdef
+  result = subprocess.call(['sh', headers_install, h, out_h], env=env)
 
   if result != 0:
     print('error: run_headers_install: cmd %s failed %d' % (cmd, result))
@@ -510,6 +513,7 @@
 
   # Tools and tool files.
   headers_install_sh = 'headers_install.sh'
+  unifdef = 'unifdef'
   kernel_headers_py = 'kernel_headers.py'
   arm_syscall_tool = 'arch/arm/tools/syscallhdr.sh'
 
@@ -657,7 +661,10 @@
 
     f.write('genrule {\n')
     f.write('    name: "qti_generate_kernel_headers_%s",\n' % header_arch)
-    f.write('    tools: ["%s"],\n' % headers_install_sh)
+    f.write('    tools: [\n')
+    f.write('        "%s",\n' % headers_install_sh)
+    f.write('        "%s",\n' % unifdef)
+    f.write('    ],\n')
     f.write('    tool_files: [\n')
     f.write('        "%s",\n' % kernel_headers_py)
 
@@ -691,6 +698,7 @@
       f.write('        "--arch_syscall_tbl $(location %s) " +\n' % arm_syscall_tbl)
 
     f.write('        "--headers_install $(location %s) " +\n' % headers_install_sh)
+    f.write('        "--unifdef $(location %s) " +\n' % unifdef)
     f.write('        "--include_uapi $(locations %s)",\n' % generic_src)
     f.write('    out: ["linux/version.h"] + gen_headers_out_%s,\n' % header_arch)
     f.write('}\n')
@@ -745,7 +753,7 @@
 def gen_headers(
     verbose, header_arch, gen_dir, arch_asm_kbuild, asm_generic_kbuild, module_dir,
     old_gen_headers_bp, new_gen_headers_bp, version_makefile,
-    arch_syscall_tool, arch_syscall_tbl, headers_install, include_uapi,
+    arch_syscall_tool, arch_syscall_tbl, headers_install, unifdef, include_uapi,
     arch_include_uapi, techpack_include_uapi):
   """Generate the kernel headers.
 
@@ -767,6 +775,7 @@
     arch_syscall_tool: The arch script that generates syscall headers.
     arch_syscall_tbl: The arch script that defines syscall vectors.
     headers_install: The headers_install tool to process input headers.
+    unifdef: The unifdef tool used by headers_install.
     include_uapi: The list of include/uapi header files.
     arch_include_uapi: The list of arch/<arch>/include/uapi header files.
   Return:
@@ -794,20 +803,20 @@
 
   for h in include_uapi:
     if not run_headers_install(
-        verbose, gen_dir, headers_install,
+        verbose, gen_dir, headers_install, unifdef,
         uapi_include_prefix, h):
       error_count += 1
 
   for h in arch_include_uapi:
     if not run_headers_install(
-        verbose, gen_dir, headers_install,
+        verbose, gen_dir, headers_install, unifdef,
         arch_uapi_include_prefix, h):
       error_count += 1
 
   for h in techpack_include_uapi:
     techpack_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,
         techpack_uapi_include_prefix, h):
       error_count += 1
 
@@ -935,6 +944,10 @@
       required=True,
       help='The headers_install tool to process input headers.')
   parser_headers.add_argument(
+      '--unifdef',
+      required=True,
+      help='The unifdef tool used by headers_install.')
+  parser_headers.add_argument(
       '--include_uapi',
       required=True,
       nargs='*',
@@ -982,12 +995,14 @@
       print('arch_syscall_tool [%s]' % args.arch_syscall_tool)
       print('arch_syscall_tbl [%s]' % args.arch_syscall_tbl)
       print('headers_install [%s]' % args.headers_install)
+      print('unifdef [%s]' % args.unifdef)
 
     return gen_headers(
         args.verbose, args.header_arch, args.gen_dir, args.arch_asm_kbuild,
         args.asm_generic_kbuild, module_dir, args.old_gen_headers_bp, args.new_gen_headers_bp,
         args.version_makefile, args.arch_syscall_tool, args.arch_syscall_tbl,
-        args.headers_install, args.include_uapi, args.arch_include_uapi, args.techpack_include_uapi)
+        args.headers_install, args.unifdef, args.include_uapi, args.arch_include_uapi,
+        args.techpack_include_uapi)
 
   print('error: unknown mode: %s' % args.mode)
   return 1