diff --git a/heat_map.py b/heat_map.py
deleted file mode 100755
index 1009eca..0000000
--- a/heat_map.py
+++ /dev/null
@@ -1,144 +0,0 @@
-#!/usr/bin/env python2
-# -*- coding: utf-8 -*-
-# Copyright 2015 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-"""Wrapper to generate heat maps for chrome."""
-
-from __future__ import print_function
-
-import argparse
-import shutil
-import os
-import sys
-import tempfile
-
-from cros_utils import command_executer
-import heatmap_generator
-
-
-def IsARepoRoot(directory):
-  """Returns True if directory is the root of a repo checkout."""
-  return os.path.exists(os.path.join(directory, '.repo'))
-
-
-class HeatMapProducer(object):
-  """Class to produce heat map."""
-
-  def __init__(self, chromeos_root, perf_data, page_size, binary, title):
-    self.chromeos_root = os.path.realpath(chromeos_root)
-    self.perf_data = os.path.realpath(perf_data)
-    self.page_size = page_size
-    self.dir = os.path.dirname(os.path.realpath(__file__))
-    self.binary = binary
-    self.tempDir = ''
-    self.ce = command_executer.GetCommandExecuter()
-    self.loading_address = None
-    self.temp_perf = ''
-    self.temp_perf_inchroot = ''
-    self.perf_report = ''
-    self.title = title
-
-  def copyFileToChroot(self):
-    self.tempDir = tempfile.mkdtemp(
-        prefix=os.path.join(self.chromeos_root, 'src/'))
-    self.temp_perf = os.path.join(self.tempDir, 'perf.data')
-    shutil.copy2(self.perf_data, self.temp_perf)
-    self.temp_perf_inchroot = os.path.join('~/trunk/src',
-                                           os.path.basename(self.tempDir))
-
-  def getPerfReport(self):
-    if os.path.isfile(os.path.join(self.tempDir, 'perf_report.txt')):
-      self.perf_report = os.path.join(self.tempDir, 'perf_report.txt')
-      return
-
-    cmd = ('cd %s; perf report -D -i perf.data > perf_report.txt' %
-           self.temp_perf_inchroot)
-    retval = self.ce.ChrootRunCommand(self.chromeos_root, cmd)
-    if retval:
-      raise RuntimeError('Failed to generate perf report')
-    self.perf_report = os.path.join(self.tempDir, 'perf_report.txt')
-
-  def getHeatMap(self, top_n_pages=None):
-    generator = heatmap_generator.HeatmapGenerator(self.perf_report,
-                                                   self.page_size, self.title)
-    generator.draw()
-    # Analyze top N hottest symbols with the binary, if provided
-    if self.binary:
-      if top_n_pages is not None:
-        generator.analyze(self.binary, top_n_pages)
-      else:
-        generator.analyze(self.binary)
-
-  def RemoveFiles(self):
-    shutil.rmtree(self.tempDir)
-    if os.path.isfile(os.path.join(os.getcwd(), 'out.txt')):
-      os.remove(os.path.join(os.getcwd(), 'out.txt'))
-    if os.path.isfile(os.path.join(os.getcwd(), 'inst-histo.txt')):
-      os.remove(os.path.join(os.getcwd(), 'inst-histo.txt'))
-
-
-def main(argv):
-  """Parse the options.
-
-  Args:
-    argv: The options with which this script was invoked.
-
-  Returns:
-    0 unless an exception is raised.
-  """
-  parser = argparse.ArgumentParser()
-
-  parser.add_argument(
-      '--chromeos_root',
-      dest='chromeos_root',
-      required=True,
-      help='ChromeOS root to use for generate heatmaps.')
-  parser.add_argument(
-      '--perf_data', dest='perf_data', required=True, help='The raw perf data.')
-  parser.add_argument(
-      '--binary',
-      dest='binary',
-      required=False,
-      help='The path to the Chrome binary.',
-      default=None)
-  parser.add_argument(
-      '--top_n',
-      dest='top_n',
-      required=False,
-      help='Print out top N hottest pages within/outside huge page range(30MB)',
-      default=None)
-  parser.add_argument(
-      '--page_size',
-      dest='page_size',
-      required=False,
-      help='The page size for heat maps.',
-      default=4096)
-  parser.add_argument('--title', dest='title', default='')
-
-  options = parser.parse_args(argv)
-
-  if not IsARepoRoot(options.chromeos_root):
-    parser.error('% does not contain .repo dir.' % options.chromeos_root)
-
-  if not os.path.isfile(options.perf_data):
-    parser.error('Cannot find perf_data: %s.' % options.perf_data)
-
-  heatmap_producer = HeatMapProducer(options.chromeos_root, options.perf_data,
-                                     options.page_size, options.binary,
-                                     options.title)
-  try:
-    heatmap_producer.copyFileToChroot()
-    heatmap_producer.getPerfReport()
-    heatmap_producer.getHeatMap(options.top_n)
-    print('\nheat map and time histgram genereated in the current directory '
-          'with name heat_map.png and timeline.png accordingly.')
-  except RuntimeError, e:
-    print(e)
-  finally:
-    heatmap_producer.RemoveFiles()
-
-
-if __name__ == '__main__':
-  sys.exit(main(sys.argv[1:]))
diff --git a/heatmaps/heat_map.py b/heatmaps/heat_map.py
new file mode 100755
index 0000000..b7006e3
--- /dev/null
+++ b/heatmaps/heat_map.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Wrapper to generate heat maps for chrome."""
+
+from __future__ import print_function
+
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+
+from cros_utils import command_executer
+import heatmap_generator
+
+
+def IsARepoRoot(directory):
+  """Returns True if directory is the root of a repo checkout."""
+  return os.path.exists(
+      os.path.join(os.path.realpath(os.path.expanduser(directory)), '.repo'))
+
+
+class HeatMapProducer(object):
+  """Class to produce heat map."""
+
+  def __init__(self,
+               chromeos_root,
+               perf_data,
+               hugepage,
+               binary,
+               title,
+               logger=None):
+    self.chromeos_root = os.path.realpath(os.path.expanduser(chromeos_root))
+    self.perf_data = os.path.realpath(os.path.expanduser(perf_data))
+    self.hugepage = hugepage
+    self.dir = os.path.dirname(os.path.realpath(__file__))
+    self.binary = binary
+    self.ce = command_executer.GetCommandExecuter()
+    self.temp_dir = ''
+    self.temp_perf_inchroot = ''
+    self.temp_dir_created = False
+    self.perf_report = ''
+    self.title = title
+    self.logger = logger
+
+  def _EnsureFileInChroot(self):
+    chroot_prefix = os.path.join(self.chromeos_root, 'chroot')
+    if self.perf_data.startswith(chroot_prefix):
+      # If the path to perf_data starts with the same chromeos_root, assume
+      # it's in the chromeos_root so no need for temporary directory and copy.
+      self.temp_dir = self.perf_data.replace('perf.data', '')
+      self.temp_perf_inchroot = self.temp_dir.replace(chroot_prefix, '')
+
+    else:
+      # Otherwise, create a temporary directory and copy perf.data into chroot.
+      self.temp_dir = tempfile.mkdtemp(
+          prefix=os.path.join(self.chromeos_root, 'src/'))
+      temp_perf = os.path.join(self.temp_dir, 'perf.data')
+      shutil.copy2(self.perf_data, temp_perf)
+      self.temp_perf_inchroot = os.path.join('~/trunk/src',
+                                             os.path.basename(self.temp_dir))
+      self.temp_dir_created = True
+
+  def _GeneratePerfReport(self):
+    cmd = ('cd %s && perf report -D -i perf.data > perf_report.txt' %
+           self.temp_perf_inchroot)
+    retval = self.ce.ChrootRunCommand(self.chromeos_root, cmd)
+    if retval:
+      raise RuntimeError('Failed to generate perf report')
+    self.perf_report = os.path.join(self.temp_dir, 'perf_report.txt')
+
+  def _GetHeatMap(self, top_n_pages):
+    generator = heatmap_generator.HeatmapGenerator(
+        perf_report=self.perf_report,
+        page_size=4096,
+        hugepage=self.hugepage,
+        title=self.title)
+    generator.draw()
+    # Analyze top N hottest symbols with the binary, if provided
+    if self.binary:
+      if top_n_pages is not None:
+        generator.analyze(self.binary, top_n_pages)
+      else:
+        generator.analyze(self.binary)
+
+  def _RemoveFiles(self):
+    files = [
+        'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt', 'inst-histo-sp.txt'
+    ]
+    for f in files:
+      if os.path.exists(f):
+        os.remove(f)
+
+  def Run(self, top_n_pages=None):
+    try:
+      self._EnsureFileInChroot()
+      self._GeneratePerfReport()
+      self._GetHeatMap(top_n_pages)
+    finally:
+      self._RemoveFiles()
+    msg = ('heat map and time histogram genereated in the current '
+           'directory with name heat_map.png and timeline.png '
+           'accordingly.')
+    if self.logger:
+      self.logger.LogOutput(msg)
+    else:
+      print(msg)
+
+
+def main(argv):
+  """Parse the options.
+
+  Args:
+    argv: The options with which this script was invoked.
+
+  Returns:
+    0 unless an exception is raised.
+  """
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument(
+      '--chromeos_root',
+      dest='chromeos_root',
+      required=True,
+      help='ChromeOS root to use for generate heatmaps.')
+  parser.add_argument(
+      '--perf_data',
+      dest='perf_data',
+      required=True,
+      help='The raw perf data. Must be collected with -e instructions while '
+      'disabling ASLR.')
+  parser.add_argument(
+      '--binary',
+      dest='binary',
+      help='The path to the Chrome binary. Only useful if want to print '
+      'symbols on hottest pages',
+      default=None)
+  parser.add_argument(
+      '--top_n',
+      dest='top_n',
+      help='Print out top N hottest pages within/outside huge page range. '
+      'Must be used with --hugepage and --binary.',
+      default=None)
+  parser.add_argument(
+      '--title', dest='title', help='Title of the heatmap', default='')
+  parser.add_argument(
+      '--hugepage',
+      dest='hugepage',
+      help='A range of addresses (start,end) where huge page starts and ends'
+      ' in text section, separated by a comma.'
+      ' Used to differentiate regions in heatmap.'
+      ' Example: --hugepage=0,4096'
+      ' If not specified, no effect on the heatmap.',
+      default=None)
+
+  options = parser.parse_args(argv)
+
+  if not IsARepoRoot(options.chromeos_root):
+    parser.error('%s does not contain .repo dir.' % options.chromeos_root)
+
+  if not os.path.isfile(options.perf_data):
+    parser.error('Cannot find perf_data: %s.' % options.perf_data)
+
+  hugepage_range = None
+  if options.hugepage:
+    hugepage_range = options.hugepage.split(',')
+    if len(hugepage_range) != 2 or \
+     int(hugepage_range[0]) > int(hugepage_range[1]):
+      parser.error('Wrong format of hugepage range: %s' % options.hugepage)
+    hugepage_range = [int(x) for x in hugepage_range]
+
+  heatmap_producer = HeatMapProducer(options.chromeos_root, options.perf_data,
+                                     hugepage_range, options.binary,
+                                     options.title)
+
+  heatmap_producer.Run(options.top_n)
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv[1:]))
diff --git a/heatmaps/heat_map_test.py b/heatmaps/heat_map_test.py
new file mode 100755
index 0000000..c647467
--- /dev/null
+++ b/heatmaps/heat_map_test.py
@@ -0,0 +1,157 @@
+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for heat_map.py."""
+
+from __future__ import print_function
+
+import mock
+import unittest
+
+import os
+
+from cros_utils import command_executer
+
+import heat_map
+
+
+def make_heatmap(chromeos_root='/path/to/fake/chromeos_root/',
+                 perf_data='/any_path/perf.data'):
+  return heat_map.HeatMapProducer(chromeos_root, perf_data, None, None, '')
+
+
+def fake_mkdtemp(prefix):
+  """Mock tempfile.mkdtemp() by just create a pathname."""
+  return prefix + 'random_dir'
+
+
+def fake_parser_error(_, msg):
+  """Redirect parser.error() to exception."""
+  raise Exception(msg)
+
+
+def fake_generate_perf_report_exception(_):
+  raise Exception
+
+
+class HeatmapTest(unittest.TestCase):
+  """All of our tests for heat_map."""
+
+  #pylint: disable=protected-access
+  @mock.patch('shutil.copy2')
+  @mock.patch('tempfile.mkdtemp')
+  def test_EnsureFileInChrootAlreadyInside(self, mock_mkdtemp, mock_copy):
+    perf_data_inchroot = (
+        '/path/to/fake/chromeos_root/chroot/inchroot_path/perf.data')
+    heatmap = make_heatmap(perf_data=perf_data_inchroot)
+    heatmap._EnsureFileInChroot()
+    self.assertFalse(heatmap.temp_dir_created)
+    self.assertEqual(heatmap.temp_dir,
+                     '/path/to/fake/chromeos_root/chroot/inchroot_path/')
+    self.assertEqual(heatmap.temp_perf_inchroot, '/inchroot_path/')
+    mock_mkdtemp.assert_not_called()
+    mock_copy.assert_not_called()
+
+  @mock.patch('shutil.copy2')
+  @mock.patch('tempfile.mkdtemp', fake_mkdtemp)
+  def test_EnsureFileInChrootOutsideNeedCopy(self, mock_copy):
+    heatmap = make_heatmap()
+    heatmap._EnsureFileInChroot()
+    self.assertTrue(heatmap.temp_dir_created)
+    self.assertEqual(mock_copy.call_count, 1)
+    self.assertEqual(heatmap.temp_dir,
+                     '/path/to/fake/chromeos_root/src/random_dir')
+    self.assertEqual(heatmap.temp_perf_inchroot, '~/trunk/src/random_dir')
+
+  @mock.patch.object(command_executer.CommandExecuter, 'ChrootRunCommand')
+  def test_GeneratePerfReport(self, mock_ChrootRunCommand):
+    heatmap = make_heatmap()
+    heatmap.temp_dir = '/fake/chroot/inchroot_path/'
+    heatmap.temp_perf_inchroot = '/inchroot_path/'
+    mock_ChrootRunCommand.return_value = 0
+    heatmap._GeneratePerfReport()
+    cmd = ('cd %s && perf report -D -i perf.data > perf_report.txt' %
+           heatmap.temp_perf_inchroot)
+    mock_ChrootRunCommand.assert_called_with(heatmap.chromeos_root, cmd)
+    self.assertEqual(mock_ChrootRunCommand.call_count, 1)
+    self.assertEqual(heatmap.perf_report,
+                     '/fake/chroot/inchroot_path/perf_report.txt')
+
+  @mock.patch('heatmap_generator.HeatmapGenerator')
+  def test_GetHeatMap(self, mock_heatmap_generator):
+    heatmap = make_heatmap()
+    heatmap._GetHeatMap(None)
+    self.assertTrue(mock_heatmap_generator.called)
+
+  @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
+  @mock.patch.object(heat_map.HeatMapProducer, '_GeneratePerfReport')
+  @mock.patch.object(heat_map.HeatMapProducer, '_GetHeatMap')
+  @mock.patch.object(heat_map.HeatMapProducer, '_RemoveFiles')
+  def test_Run(self, mock_remove_files, mock_get_heatmap,
+               mock_generate_perf_report, mock_ensure_file_in_chroot):
+    heatmap = make_heatmap()
+    heatmap.Run()
+    mock_ensure_file_in_chroot.assert_called_once_with()
+    mock_generate_perf_report.assert_called_once_with()
+    mock_get_heatmap.assert_called_once_with(None)
+    mock_remove_files.assert_called_once_with()
+
+  @mock.patch.object(heat_map.HeatMapProducer, '_EnsureFileInChroot')
+  @mock.patch.object(
+      heat_map.HeatMapProducer,
+      '_GeneratePerfReport',
+      new=fake_generate_perf_report_exception)
+  @mock.patch.object(heat_map.HeatMapProducer, '_GetHeatMap')
+  @mock.patch.object(heat_map.HeatMapProducer, '_RemoveFiles')
+  @mock.patch('__builtin__.print')
+  def test_Run_with_exception(self, mock_print, mock_remove_files,
+                              mock_get_heatmap, mock_ensure_file_in_chroot):
+    heatmap = make_heatmap()
+    with self.assertRaises(Exception):
+      heatmap.Run()
+    mock_ensure_file_in_chroot.assert_called_once_with()
+    mock_get_heatmap.assert_not_called()
+    mock_remove_files.assert_called_once_with()
+    mock_print.assert_not_called()
+
+  @mock.patch('argparse.ArgumentParser.error', fake_parser_error)
+  @mock.patch.object(os.path, 'isfile')
+  @mock.patch.object(heat_map, 'IsARepoRoot')
+  def test_main_arg_format(self, mock_IsARepoRoot, mock_isfile):
+    """Test wrong arg format are detected."""
+    args = ['--chromeos_root=/fake/chroot/', '--perf_data=/path/to/perf.data']
+
+    # Test --chromeos_root format
+    mock_IsARepoRoot.return_value = False
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('does not contain .repo dir.', str(msg.exception))
+
+    # Test --perf_data format
+    mock_IsARepoRoot.return_value = True
+    mock_isfile.return_value = False
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('Cannot find perf_data', str(msg.exception))
+
+    # Test --hugepage format
+    mock_isfile.return_value = True
+    args.append('--hugepage=0')
+    with self.assertRaises(Exception) as msg:
+      heat_map.main(args)
+    self.assertIn('Wrong format of hugepage range', str(msg.exception))
+
+    # Test --hugepage parse
+    args[-1] = '--hugepage=0,4096'
+    heat_map.HeatMapProducer = mock.MagicMock()
+    heat_map.main(args)
+    heat_map.HeatMapProducer.assert_called_with(
+        '/fake/chroot/', '/path/to/perf.data', [0, 4096], None, '')
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/heatmap_generator.py b/heatmaps/heatmap_generator.py
similarity index 93%
rename from heatmap_generator.py
rename to heatmaps/heatmap_generator.py
index 2b7bf92..4f76450 100644
--- a/heatmap_generator.py
+++ b/heatmaps/heatmap_generator.py
@@ -23,6 +23,8 @@
 
 from cros_utils import command_executer
 
+HugepageRange = collections.namedtuple('HugepageRange', ['start', 'end'])
+
 
 class MMap(object):
   """Class to store mmap information in perf report.
@@ -78,7 +80,12 @@
   chrome binary
   """
 
-  def __init__(self, perf_report, page_size, title, log_level='verbose'):
+  def __init__(self,
+               perf_report,
+               page_size,
+               hugepage,
+               title,
+               log_level='verbose'):
     self.perf_report = perf_report
     # Pick 1G as a relatively large number. All addresses less than it will
     # be recorded. The actual heatmap will show up to a boundary of the
@@ -93,13 +100,14 @@
     self.processes = {}
     self.deleted_processes = {}
     self.count = 0
-    self.page_size = page_size
+    if hugepage:
+      self.hugepage = HugepageRange(start=hugepage[0], end=hugepage[1])
+    else:
+      self.hugepage = None
     self.title = title
     self.symbol_addresses = []
     self.symbol_names = []
-    # Set huge page region of Chrome to be first 30MB, only used when printing
-    # out hottest pages
-    self.huge_page_size = 30 * 1024 * 1024
+    self.page_size = page_size
 
   def _parse_perf_sample(self, line):
     # In a perf report, generated with -D, a PERF_RECORD_SAMPLE command should
@@ -268,9 +276,14 @@
     ' the starting address of Chrome'
     if address < self.max_addr:
       self.count += 1
-      print(('%d/%d: %d %d' % (pid[0], pid[1], self.count,
-                               address / self.page_size * self.page_size)),
-            file=self.hist_temp_output)
+      line = '%d/%d: %d %d' % (pid[0], pid[1], self.count,
+                               address / self.page_size * self.page_size)
+      if self.hugepage:
+        if self.hugepage.start <= address < self.hugepage.end:
+          line += ' hugepage'
+        else:
+          line += ' smallpage'
+      print(line, file=self.hist_temp_output)
 
   def _read_perf_report(self):
     # Serve as main function to read perf report, generated by -D
@@ -292,7 +305,12 @@
     # of results written in out.txt and also generate pngs for
     # heat maps.
     heatmap_script = os.path.join(self.dir, 'perf-to-inst-page.sh')
-    cmd = '{0} {1}'.format(heatmap_script, pipes.quote(self.title))
+    if self.hugepage:
+      hp_arg = 'hugepage'
+    else:
+      hp_arg = 'none'
+
+    cmd = '{0} {1} {2}'.format(heatmap_script, pipes.quote(self.title), hp_arg)
     retval = self.ce.RunCommand(cmd)
     if retval:
       raise RuntimeError('Failed to run script to generate heatmap')
@@ -366,11 +384,11 @@
       list_to_show = [(k, v) for (k, v) in sorted_hist if v >= count_threshold]
     else:
       in_huge_page = [
-          (k, v) for (k, v) in sorted_hist if k < self.huge_page_size
-      ][:int(top_n)]
+          (k, v) for (k, v) in sorted_hist \
+          if self.hugepage.start <= k < self.hugepage.end][:int(top_n)]
       outside_huge_page = [
-          (k, v) for (k, v) in sorted_hist if k >= self.huge_page_size
-      ][:int(top_n)]
+          (k, v) for (k, v) in sorted_hist \
+          if k < self.hugepage.start or k >= self.hugepage.end][:int(top_n)]
 
       list_to_show = in_huge_page + outside_huge_page
     return list_to_show, max_count
diff --git a/heatmap_generator_test.py b/heatmaps/heatmap_generator_test.py
similarity index 79%
rename from heatmap_generator_test.py
rename to heatmaps/heatmap_generator_test.py
index 53f7dd2..271f9f4 100755
--- a/heatmap_generator_test.py
+++ b/heatmaps/heatmap_generator_test.py
@@ -45,18 +45,21 @@
   print(' ...... dso: /opt/google/chrome/chrome\n', file=fp)
 
 
-def _heatmap(file_name, page_size=4096):
+def _heatmap(file_name, page_size=4096, hugepage=None):
   generator = heatmap_generator.HeatmapGenerator(
-      file_name, page_size, '', log_level='none')  # Don't log to stdout
+      file_name, page_size, hugepage, '',
+      log_level='none')  # Don't log to stdout
   generator.draw()
 
 
 def _cleanup(file_name):
-  os.remove(file_name)
-  os.remove('out.txt')
-  os.remove('inst-histo.txt')
-  os.remove('heat_map.png')
-  os.remove('timeline.png')
+  files = [
+      file_name, 'out.txt', 'inst-histo.txt', 'inst-histo-hp.txt',
+      'inst-histo-sp.txt', 'heat_map.png', 'timeline.png'
+  ]
+  for f in files:
+    if os.path.exists(f):
+      os.remove(f)
 
 
 class Tests(unittest.TestCase):
@@ -68,13 +71,13 @@
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
       _write_perf_sample(101, 101, 0xABCD101, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
       lines = f.readlines()
       self.assertEqual(len(lines), 1)
       self.assertIn('101/101: 1 0', lines[0])
-    _cleanup(fname)
 
   def test_with_one_mmap_multiple_samples(self):
     """Tests one perf record and three samples."""
@@ -84,6 +87,7 @@
       _write_perf_sample(101, 101, 0xABCD101, f)
       _write_perf_sample(101, 101, 0xABCD102, f)
       _write_perf_sample(101, 101, 0xABCE102, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -92,7 +96,6 @@
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('101/101: 2 0', lines[1])
       self.assertIn('101/101: 3 4096', lines[2])
-    _cleanup(fname)
 
   def test_with_fork_and_exit(self):
     """Tests perf fork and perf exit."""
@@ -103,6 +106,7 @@
       _write_perf_sample(101, 101, 0xABCD101, f)
       _write_perf_sample(202, 202, 0xABCE101, f)
       _write_perf_exit(202, 202, 202, 202, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -110,7 +114,6 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('202/202: 2 4096', lines[1])
-    _cleanup(fname)
 
   def test_hugepage_creates_two_chrome_mmaps(self):
     """Test two chrome mmaps for the same process."""
@@ -122,6 +125,7 @@
       _write_perf_mmap(202, 202, 0xABCD300, 0xD00, f)
       _write_perf_sample(101, 101, 0xABCD102, f)
       _write_perf_sample(202, 202, 0xABCD102, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('out.txt', os.listdir('.'))
     with open('out.txt') as f:
@@ -129,7 +133,6 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('101/101: 1 0', lines[0])
       self.assertIn('202/202: 2 0', lines[1])
-    _cleanup(fname)
 
   def test_hugepage_creates_two_chrome_mmaps_fail(self):
     """Test two chrome mmaps for the same process."""
@@ -139,6 +142,7 @@
       _write_perf_mmap(101, 101, 0xABCD000, 0x1000, f)
       _write_perf_fork(101, 101, 202, 202, f)
       _write_perf_mmap(202, 202, 0xABCD000, 0x10000, f)
+    self.addCleanup(_cleanup, fname)
     with self.assertRaises(AssertionError) as msg:
       _heatmap(fname)
     self.assertIn('Original MMAP size', str(msg.exception))
@@ -173,11 +177,12 @@
     fname = 'test_histo.txt'
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
-      for i in range(0, 100):
+      for i in range(100):
         _write_perf_sample(101, 101, 0xABCD000 + i, f)
         _write_perf_sample(101, 101, 0xABCE000 + i, f)
         _write_perf_sample(101, 101, 0xABFD000 + i, f)
         _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname)
     self.assertIn('inst-histo.txt', os.listdir('.'))
     with open('inst-histo.txt') as f:
@@ -187,18 +192,18 @@
       self.assertIn('100 4096', lines[1])
       self.assertIn('100 196608', lines[2])
       self.assertIn('100 4194304', lines[3])
-    _cleanup(fname)
 
   def test_histogram_two_mb_page(self):
     """Tests handling of 2MB page."""
     fname = 'test_histo.txt'
     with open(fname, 'w') as f:
       _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
-      for i in range(0, 100):
+      for i in range(100):
         _write_perf_sample(101, 101, 0xABCD000 + i, f)
         _write_perf_sample(101, 101, 0xABCE000 + i, f)
         _write_perf_sample(101, 101, 0xABFD000 + i, f)
         _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
     _heatmap(fname, page_size=2 * 1024 * 1024)
     self.assertIn('inst-histo.txt', os.listdir('.'))
     with open('inst-histo.txt') as f:
@@ -206,7 +211,33 @@
       self.assertEqual(len(lines), 2)
       self.assertIn('300 0', lines[0])
       self.assertIn('100 4194304', lines[1])
-    _cleanup(fname)
+
+  def test_histogram_in_and_out_hugepage(self):
+    """Tests handling the case of separating samples in and out huge page."""
+    fname = 'test_histo.txt'
+    with open(fname, 'w') as f:
+      _write_perf_mmap(101, 101, 0xABCD000, 0x100, f)
+      for i in range(100):
+        _write_perf_sample(101, 101, 0xABCD000 + i, f)
+        _write_perf_sample(101, 101, 0xABCE000 + i, f)
+        _write_perf_sample(101, 101, 0xABFD000 + i, f)
+        _write_perf_sample(101, 101, 0xAFCD000 + i, f)
+    self.addCleanup(_cleanup, fname)
+    _heatmap(fname, hugepage=[0, 8192])
+    file_list = os.listdir('.')
+    self.assertNotIn('inst-histo.txt', file_list)
+    self.assertIn('inst-histo-hp.txt', file_list)
+    self.assertIn('inst-histo-sp.txt', file_list)
+    with open('inst-histo-hp.txt') as f:
+      lines = f.readlines()
+      self.assertEqual(len(lines), 2)
+      self.assertIn('100 0', lines[0])
+      self.assertIn('100 4096', lines[1])
+    with open('inst-histo-sp.txt') as f:
+      lines = f.readlines()
+      self.assertEqual(len(lines), 2)
+      self.assertIn('100 196608', lines[0])
+      self.assertIn('100 4194304', lines[1])
 
 
 if __name__ == '__main__':
diff --git a/heatmaps/perf-to-inst-page.sh b/heatmaps/perf-to-inst-page.sh
new file mode 100755
index 0000000..d6acd5e
--- /dev/null
+++ b/heatmaps/perf-to-inst-page.sh
@@ -0,0 +1,68 @@
+#! /bin/bash -u
+# Copyright 2015 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# This script takes the out.txt, generated by heatmap_generator.py
+# and sorted into a heatmap data file (inst-histo.txt) and then
+# call gnu plot to draw the heat map and the time map.
+# A heat map shows the overall hotness of instructions being executed
+# while the time map shows the hotness of instruction at different time.
+#
+# Usage:
+# ./perf-to-inst-page.sh HEATMAP_TITLE
+# HEATMAP_TITLE: the title to display on the heatmap
+
+HEAT_PNG="heat_map.png"
+TIMELINE_PNG="timeline.png"
+HEATMAP_TITLE=$1
+ENABLE_HUGEPAGE=$2
+
+heatmap_command="
+set terminal png size 600,450
+set xlabel \"Instruction Virtual Address (MB)\"
+set ylabel \"Sample Occurance\"
+set grid
+
+set output \"${HEAT_PNG}\"
+set title \"${HEATMAP_TITLE}\"
+"
+
+if [[ "${ENABLE_HUGEPAGE}" = "hugepage" ]]; then
+  hugepage_hist="inst-histo-hp.txt"
+  smallpage_hist="inst-histo-sp.txt"
+  cat out.txt | grep hugepage | awk '{print $3}' \
+    | sort -n | uniq -c > "${hugepage_hist}"
+  cat out.txt | grep smallpage | awk '{print $3}' \
+    | sort -n | uniq -c > "${smallpage_hist}"
+  # generate inst heat map
+  heatmap_in_hugepage=("'${hugepage_hist}' using \
+(\$2/1024/1024):1 with impulses notitle lt rgb 'red'")
+  heatmap_outside_hugepage=("'${smallpage_hist}' using \
+(\$2/1024/1024):1 with impulses notitle lt rgb 'blue'")
+  echo "
+  ${heatmap_command}
+  plot ${heatmap_in_hugepage}, ${heatmap_outside_hugepage}
+  " | gnuplot
+else
+  awk '{print $3}' out.txt | sort -n | uniq -c > inst-histo.txt
+  # generate inst heat map
+  echo "
+  ${heatmap_command}
+  plot 'inst-histo.txt' using (\$2/1024/1024):1 with impulses notitle
+  " | gnuplot
+fi
+
+# generate instruction page access timeline
+num=$(awk 'END {print NR+1}' out.txt)
+
+echo "
+set terminal png size 600,450
+set xlabel \"time (sec)\"
+set ylabel \"Instruction Virtual Address (MB)\"
+
+set output \"${TIMELINE_PNG}\"
+set title \"instruction page accessd timeline\"
+
+plot 'out.txt' using (\$0/$num*10):(\$3/1024/1024) with dots notitle
+" | gnuplot
diff --git a/perf-to-inst-page.sh b/perf-to-inst-page.sh
deleted file mode 100755
index ea6e43c..0000000
--- a/perf-to-inst-page.sh
+++ /dev/null
@@ -1,47 +0,0 @@
-#! /bin/bash -u
-# Copyright 2015 The Chromium OS Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-# This script takes the out.txt, generated by heatmap_generator.py
-# and sorted into a heatmap data file (inst-histo.txt) and then
-# call gnu plot to draw the heat map and the time map.
-# A heat map shows the overall hotness of instructions being executed
-# while the time map shows the hotness of instruction at different time.
-#
-# Usage:
-# ./perf-to-inst-page.sh HEATMAP_TITLE
-# HEATMAP_TITLE: the title to display on the heatmap
-
-HEAT_PNG="heat_map.png"
-TIMELINE_PNG="timeline.png"
-HEATMAP_TITLE=$1
-
-awk '{print $3}' out.txt | sort -n | uniq -c > inst-histo.txt
-
-# generate inst heat map
-echo "
-set terminal png size 600,450
-set xlabel \"Instruction Virtual Address (MB)\"
-set ylabel \"Sample Occurance\"
-set grid
-
-set output \"${HEAT_PNG}\"
-set title \"${HEATMAP_TITLE}\"
-
-plot 'inst-histo.txt' using (\$2/1024/1024):1 with impulses notitle
-" | gnuplot
-
-# generate instruction page access timeline
-num=$(awk 'END {print NR+1}' out.txt)
-
-echo "
-set terminal png size 600,450
-set xlabel \"time (sec)\"
-set ylabel \"Instruction Virtual Address (MB)\"
-
-set output \"${TIMELINE_PNG}\"
-set title \"instruction page accessd timeline\"
-
-plot 'out.txt' using (\$0/$num*10):(\$3/1024/1024) with dots notitle
-" | gnuplot