tests/python/contrib/test_hexagon/test_benchmark_maxpool2d.py - tvm - Git at Google

 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
 # with the License.  You may obtain a copy of the License at
 #
 #   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
 # ruff: noqa: F401, RUF012

 """
 This module serves two purposes:
     (1) Demonstrates how to write Python code that exercises various
         Hexagon-related algorithms / features.

     (2) Benchmark the resulting primfuncs.

 Current limitations:
     - Input shapes are limited to NHWC --> NHWC_8h8w32c.

     - Testing parameters (input shapes, dtypes, etc.) currently
       support only one value for each parameter.

     - height, width, channel must be integer multiples of 8, 8, and 32,
       respectively.  I.e., partial blocks aren't currently
       supported by this script.

     - Requires that I/O tensors reside in "global.VTCM" memory,
       rather than "global" memory.
       This prevents benchmarking with I/O tensors that are too
       large to fit into availble VTCM.

     - The script only develops one primfunc.
       Future revisions to this script are expected to add more
       primfuncs and demonstrate more coding strategies.
 """

 import copy
 import os
 from typing import List

 import numpy as np
 import pytest

 import tvm.testing
 from tvm import te, tir, topi
 from tvm.contrib.hexagon import allocate_hexagon_array
 from tvm.contrib.hexagon.session import Session
 from tvm.topi import testing

 from . import benchmark_util as bu
 from .infrastructure import get_hexagon_target

 # Pytest seems to require that fixture names exist in the current module.
 # E.g., it doesn't allow: @pytest.mark.usefixtures("bu.benchmark_group")
 BENCHMARK_GROUP = bu.benchmark_group

 _SHOULD_SKIP_BENCHMARKS, _SKIP_BENCHMARKS_REASON = bu.skip_benchmarks_flag_and_reason()


 def _ceil_div(numerator, denominator):
     return (numerator + (denominator - 1)) // denominator


 def _int8_nhwc_8h8w32c_map(n_batch, height, width, channel):
     return [
         n_batch,
         height // 8,
         width // 8,
         channel // 32,
         te.AXIS_SEPARATOR,
         height % 8,
         width % 8,
         channel % 32,
     ]


 def _int8_nhwc_8h8w32c_shape(n_batch, height, width, channel) -> List[int]:
     return [
         n_batch,
         _ceil_div(height, 8),
         _ceil_div(width, 8),
         _ceil_div(channel, 32),
         8,
         8,
         32,
     ]


 def _int8_nhwc_8h8w32c_xform_immediate(arr_in: np.ndarray) -> np.ndarray:
     """
     Return a deep copy of 'arr_in', transformed from a NWHC to
     NHWC-8h8wc32 shape.  Any newly created array elements have value 0.
     """
     stage1 = copy.copy(arr_in)

     (
         n_batch,
         height,
         width,
         channel,
     ) = stage1.shape

     (
         h_minor,
         w_minor,
         c_minor,
     ) = [8, 8, 32]

     h_major = _ceil_div(height, h_minor)
     w_major = _ceil_div(width, w_minor)
     c_major = _ceil_div(channel, c_minor)

     # This handles cases where the dimensions of arr_in are not cleanly divided
     # by the minor block size, i.e. [8, 8, 32].
     #
     # Any additional array elements that this creates will ahve value 0.
     # We shouldn't actually care what value is used for those elements, because they
     # shouldn't be treated as meaningful by any of our algorithms.
     if (height % h_minor) or (width % w_minor) or (channel % c_minor):
         stage1.resize(
             (n_batch, h_major * h_minor, w_major * w_minor, c_major * c_minor), refcheck=False
         )

     stage2 = stage1.reshape(n_batch, h_major, h_minor, w_major, w_minor, c_major, c_minor)
     stage3 = stage2.transpose(0, 1, 3, 5, 2, 4, 6)
     return stage3


 def _create_test_input(shape, dtype: str) -> np.ndarray:
     np_dtype = np.dtype(dtype)
     min_value = np.iinfo(np_dtype).min
     max_value = np.iinfo(np_dtype).max
     return np.random.randint(low=min_value, high=max_value, size=tuple(shape), dtype=np.int8)


 @pytest.mark.usefixtures("BENCHMARK_GROUP")
 class TestMaxPool2D:
     """maxpool2D base test class"""

     csv_column_order = [
         # Identifies which TE-compute / TIRScript is used as the basis for the
         # benchmarked primfunc. Only needs to be meaningful to humans.
         "basic_kernel",
         # When applicable, indicates the particular variation of schedules
         # apply by the Python code. Decoding this may require looking at this
         # script's source code.
         "sched_type",
         # Values directly based on test parameters...
         "input_shape_4d",
         "block_shape",
         "dtype",
         "kernel",
         "stride",
         "dilation",
         "padding",
         "io_tensor_mem_scope",
         # Reserved columns defined by the BenchmarksTable class.
         "row_status",
         "timings_min_usecs",
         "timings_max_usecs",
         "timings_median_usecs",
         "timings_mean_usecs",
         "timings_stddev_usecs",
         # For benchmarks that produce files on the host file system, this indicates
         # their location. Useful for post-mortem investigation of benchmark results.
         "host_files_dir_path",
         # Miscellaneous comments about the benchmark.
         "comments",
     ]

     dtype = tvm.testing.parameter("int8")

     # FIXME(cconvey): The script currently fails when height, width, or channel is not an
     # integer multiple of 8, 8, or 32, respectively.
     n_batch = tvm.testing.parameter(1)
     height = tvm.testing.parameter(*[x * 8 for x in [1, 4, 16]])
     width = tvm.testing.parameter(*[x * 8 for x in [1, 4, 16]])
     channel = tvm.testing.parameter(*[x * 32 for x in [1, 2]])

     kernel = tvm.testing.parameter((1, 1), (3, 3))
     stride = tvm.testing.parameter((1, 1))
     dilation = tvm.testing.parameter((1, 1))
     padding = tvm.testing.parameter((0, 0, 0, 0))
     io_tensor_mem_scope = tvm.testing.parameter("global.vtcm")

     @pytest.mark.skipif(_SHOULD_SKIP_BENCHMARKS, reason=_SKIP_BENCHMARKS_REASON)
     @tvm.testing.requires_hexagon
     def test_maxpool2d_nhwc(
         self,
         n_batch,
         height,
         width,
         channel,
         dtype,
         kernel,
         stride,
         dilation,
         padding,
         io_tensor_mem_scope,
         hexagon_session: Session,
     ):
         """Test maxpool2d NHWC"""

         keys_dict = {
             "basic_kernel": "max_pool2d",
             "sched_type": 1,
             "input_shape_4d": [n_batch, height, width, channel],
             "block_shape": [8, 8, 32],
             "dtype": dtype,
             "kernel": kernel,
             "stride": stride,
             "dilation": dilation,
             "padding": padding,
             "io_tensor_mem_scope": io_tensor_mem_scope,
         }

         desc = bu.get_benchmark_decription(keys_dict)

         # Create the host-side directory for this benchmark run's files / logs...
         host_files_dir_name = bu.get_benchmark_id(keys_dict)
         host_files_dir_path = os.path.join(self.working_dir, host_files_dir_name)
         os.mkdir(host_files_dir_path)

         keys_dict["host_files_dir_path"] = host_files_dir_path

         log_file_path = os.path.join(host_files_dir_path, "out.txt")
         with open(log_file_path, "w") as log_file:
             print(f"CONFIGURATION: {desc}")
             log_file.write(f"CONFIGURATION: {desc}\n")

             try:
                 input_tensor_shape_4d = [n_batch, height, width, channel]
                 input_tensor_shape_7d = _int8_nhwc_8h8w32c_shape(n_batch, height, width, channel)

                 data = te.placeholder(tuple(input_tensor_shape_4d), dtype=dtype)

                 output = topi.nn.pool2d(
                     data, kernel, stride, dilation, padding, "max", layout="NHWC"
                 )
                 primfunc = te.create_prim_func([data, output])

                 sch = tvm.s_tir.Schedule(primfunc, debug_mask="all")

                 sch.transform_layout(
                     block="tensor", buffer="placeholder", index_map=_int8_nhwc_8h8w32c_map
                 )

                 built_module = tvm.compile(
                     sch.mod,
                     target=get_hexagon_target("v69"),
                 )

                 # Save a local copy of the Hexagon object code (in the form of a .so file)
                 # to allow post-mortem inspection.
                 host_dso_binary_path = os.path.join(host_files_dir_path, "test_binary.so")
                 built_module.write_to_file(host_dso_binary_path)
                 print(f"SAVED BINARY TO HOST PATH: {host_dso_binary_path}")

                 hexagon_mod = hexagon_session.load_module(built_module)

                 # Generate the input tensor's data.
                 # Note that we'll eventually need it in two different layouts:
                 # (1) NHWC as an argument to testing.poolnd_python.
                 # (2) NHWC_8h8w32c for as an argument to our Hexagon primfunc.
                 # a_numpy_4d = np.random.randint(low=-128, high=127,
                 #   size=input_tensor_shape_4d, dtype=np.int8)
                 a_numpy_4d = _create_test_input(input_tensor_shape_4d, dtype)

                 ref_output_4d = testing.poolnd_python(
                     a_numpy_4d.astype("int32"),
                     kernel,
                     stride,
                     dilation,
                     padding[0:2],
                     padding[2:],
                     pool_type="max",
                     dtype="int32",
                     layout="NHWC",
                 ).astype(dtype)

                 output_tensor_shape_4d = ref_output_4d.shape

                 a_numpy_7d = _int8_nhwc_8h8w32c_xform_immediate(a_numpy_4d)

                 a_hexagon_7d = allocate_hexagon_array(
                     hexagon_session.device,
                     tensor_shape=input_tensor_shape_7d,
                     axis_separators=[4],
                     dtype=dtype,
                     mem_scope=io_tensor_mem_scope,
                 )

                 c_hexagon_4d = allocate_hexagon_array(
                     hexagon_session.device,
                     tensor_shape=output_tensor_shape_4d,
                     axis_separators=[],
                     dtype=dtype,
                     mem_scope=io_tensor_mem_scope,
                 )

                 a_hexagon_7d.copyfrom(a_numpy_7d)

                 if dtype == "int8":
                     rel_tolerance = 0
                     abs_tolerance = 0
                 else:
                     assert False, f"TODO: decide acceptable tolerances for dtype {dtype}"

                 timer = hexagon_mod.time_evaluator(
                     "main", hexagon_session.device, number=10, repeat=1
                 )
                 timing_result = timer(a_hexagon_7d, c_hexagon_4d)

                 try:
                     tvm.testing.assert_allclose(
                         ref_output_4d, c_hexagon_4d.numpy(), rtol=rel_tolerance, atol=abs_tolerance
                     )
                 except AssertionError as exception:
                     raise bu.NumericalAccuracyException(str(exception))

             except bu.NumericalAccuracyException as exception:
                 print()
                 print("FAIL: Numerical accuracy error. See log file.")

                 log_file.write("\n")
                 log_file.write(f"FAIL: {exception}\n")

                 self.benchmark_table.record_fail(
                     **keys_dict, comments="Numerical accuracy error. See log file."
                 )

             except bu.UnsupportedException as exception:
                 print()
                 print(f"SKIP: {exception}")

                 log_file.write("\n")
                 log_file.write(f"SKIP: {exception}\n")

                 self.benchmark_table.record_skip(
                     **keys_dict, comments=f"Unsupported configuration: {exception}"
                 )

             self.benchmark_table.record_success(timing_result, **keys_dict)


 if __name__ == "__main__":
     tvm.testing.main()
	# Licensed to the Apache Software Foundation (ASF) under one
	# or more contributor license agreements. See the NOTICE file
	# distributed with this work for additional information
	# regarding copyright ownership. The ASF licenses this file
	# to you under the Apache License, Version 2.0 (the
	# "License"); you may not use this file except in compliance
	# with the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing,
	# software distributed under the License is distributed on an
	# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	# KIND, either express or implied. See the License for the
	# specific language governing permissions and limitations
	# under the License.
	# ruff: noqa: F401, RUF012

	"""
	This module serves two purposes:
	(1) Demonstrates how to write Python code that exercises various
	Hexagon-related algorithms / features.

	(2) Benchmark the resulting primfuncs.

	Current limitations:
	- Input shapes are limited to NHWC --> NHWC_8h8w32c.

	- Testing parameters (input shapes, dtypes, etc.) currently
	support only one value for each parameter.

	- height, width, channel must be integer multiples of 8, 8, and 32,
	respectively. I.e., partial blocks aren't currently
	supported by this script.

	- Requires that I/O tensors reside in "global.VTCM" memory,
	rather than "global" memory.
	This prevents benchmarking with I/O tensors that are too
	large to fit into availble VTCM.

	- The script only develops one primfunc.
	Future revisions to this script are expected to add more
	primfuncs and demonstrate more coding strategies.
	"""

	import copy
	import os
	from typing import List

	import numpy as np
	import pytest

	import tvm.testing
	from tvm import te, tir, topi
	from tvm.contrib.hexagon import allocate_hexagon_array
	from tvm.contrib.hexagon.session import Session
	from tvm.topi import testing

	from . import benchmark_util as bu
	from .infrastructure import get_hexagon_target

	# Pytest seems to require that fixture names exist in the current module.
	# E.g., it doesn't allow: @pytest.mark.usefixtures("bu.benchmark_group")
	BENCHMARK_GROUP = bu.benchmark_group

	_SHOULD_SKIP_BENCHMARKS, _SKIP_BENCHMARKS_REASON = bu.skip_benchmarks_flag_and_reason()


	def _ceil_div(numerator, denominator):
	return (numerator + (denominator - 1)) // denominator


	def _int8_nhwc_8h8w32c_map(n_batch, height, width, channel):
	return [
	n_batch,
	height // 8,
	width // 8,
	channel // 32,
	te.AXIS_SEPARATOR,
	height % 8,
	width % 8,
	channel % 32,
	]


	def _int8_nhwc_8h8w32c_shape(n_batch, height, width, channel) -> List[int]:
	return [
	n_batch,
	_ceil_div(height, 8),
	_ceil_div(width, 8),
	_ceil_div(channel, 32),
	8,
	8,
	32,
	]


	def _int8_nhwc_8h8w32c_xform_immediate(arr_in: np.ndarray) -> np.ndarray:
	"""
	Return a deep copy of 'arr_in', transformed from a NWHC to
	NHWC-8h8wc32 shape. Any newly created array elements have value 0.
	"""
	stage1 = copy.copy(arr_in)

	(
	n_batch,
	height,
	width,
	channel,
	) = stage1.shape

	(
	h_minor,
	w_minor,
	c_minor,
	) = [8, 8, 32]

	h_major = _ceil_div(height, h_minor)
	w_major = _ceil_div(width, w_minor)
	c_major = _ceil_div(channel, c_minor)

	# This handles cases where the dimensions of arr_in are not cleanly divided
	# by the minor block size, i.e. [8, 8, 32].
	#
	# Any additional array elements that this creates will ahve value 0.
	# We shouldn't actually care what value is used for those elements, because they
	# shouldn't be treated as meaningful by any of our algorithms.
	if (height % h_minor) or (width % w_minor) or (channel % c_minor):
	stage1.resize(
	(n_batch, h_major * h_minor, w_major * w_minor, c_major * c_minor), refcheck=False
	)

	stage2 = stage1.reshape(n_batch, h_major, h_minor, w_major, w_minor, c_major, c_minor)
	stage3 = stage2.transpose(0, 1, 3, 5, 2, 4, 6)
	return stage3


	def _create_test_input(shape, dtype: str) -> np.ndarray:
	np_dtype = np.dtype(dtype)
	min_value = np.iinfo(np_dtype).min
	max_value = np.iinfo(np_dtype).max
	return np.random.randint(low=min_value, high=max_value, size=tuple(shape), dtype=np.int8)


	@pytest.mark.usefixtures("BENCHMARK_GROUP")
	class TestMaxPool2D:
	"""maxpool2D base test class"""

	csv_column_order = [
	# Identifies which TE-compute / TIRScript is used as the basis for the
	# benchmarked primfunc. Only needs to be meaningful to humans.
	"basic_kernel",
	# When applicable, indicates the particular variation of schedules
	# apply by the Python code. Decoding this may require looking at this
	# script's source code.
	"sched_type",
	# Values directly based on test parameters...
	"input_shape_4d",
	"block_shape",
	"dtype",
	"kernel",
	"stride",
	"dilation",
	"padding",
	"io_tensor_mem_scope",
	# Reserved columns defined by the BenchmarksTable class.
	"row_status",
	"timings_min_usecs",
	"timings_max_usecs",
	"timings_median_usecs",
	"timings_mean_usecs",
	"timings_stddev_usecs",
	# For benchmarks that produce files on the host file system, this indicates
	# their location. Useful for post-mortem investigation of benchmark results.
	"host_files_dir_path",
	# Miscellaneous comments about the benchmark.
	"comments",
	]

	dtype = tvm.testing.parameter("int8")

	# FIXME(cconvey): The script currently fails when height, width, or channel is not an
	# integer multiple of 8, 8, or 32, respectively.
	n_batch = tvm.testing.parameter(1)
	height = tvm.testing.parameter([x 8 for x in [1, 4, 16]])
	width = tvm.testing.parameter([x 8 for x in [1, 4, 16]])
	channel = tvm.testing.parameter([x 32 for x in [1, 2]])

	kernel = tvm.testing.parameter((1, 1), (3, 3))
	stride = tvm.testing.parameter((1, 1))
	dilation = tvm.testing.parameter((1, 1))
	padding = tvm.testing.parameter((0, 0, 0, 0))
	io_tensor_mem_scope = tvm.testing.parameter("global.vtcm")

	@pytest.mark.skipif(_SHOULD_SKIP_BENCHMARKS, reason=_SKIP_BENCHMARKS_REASON)
	@tvm.testing.requires_hexagon
	def test_maxpool2d_nhwc(
	self,
	n_batch,
	height,
	width,
	channel,
	dtype,
	kernel,
	stride,
	dilation,
	padding,
	io_tensor_mem_scope,
	hexagon_session: Session,
	):
	"""Test maxpool2d NHWC"""

	keys_dict = {
	"basic_kernel": "max_pool2d",
	"sched_type": 1,
	"input_shape_4d": [n_batch, height, width, channel],
	"block_shape": [8, 8, 32],
	"dtype": dtype,
	"kernel": kernel,
	"stride": stride,
	"dilation": dilation,
	"padding": padding,
	"io_tensor_mem_scope": io_tensor_mem_scope,
	}

	desc = bu.get_benchmark_decription(keys_dict)

	# Create the host-side directory for this benchmark run's files / logs...
	host_files_dir_name = bu.get_benchmark_id(keys_dict)
	host_files_dir_path = os.path.join(self.working_dir, host_files_dir_name)
	os.mkdir(host_files_dir_path)

	keys_dict["host_files_dir_path"] = host_files_dir_path

	log_file_path = os.path.join(host_files_dir_path, "out.txt")
	with open(log_file_path, "w") as log_file:
	print(f"CONFIGURATION: {desc}")
	log_file.write(f"CONFIGURATION: {desc}\n")

	try:
	input_tensor_shape_4d = [n_batch, height, width, channel]
	input_tensor_shape_7d = _int8_nhwc_8h8w32c_shape(n_batch, height, width, channel)

	data = te.placeholder(tuple(input_tensor_shape_4d), dtype=dtype)

	output = topi.nn.pool2d(
	data, kernel, stride, dilation, padding, "max", layout="NHWC"
	)
	primfunc = te.create_prim_func([data, output])

	sch = tvm.s_tir.Schedule(primfunc, debug_mask="all")

	sch.transform_layout(
	block="tensor", buffer="placeholder", index_map=_int8_nhwc_8h8w32c_map
	)

	built_module = tvm.compile(
	sch.mod,
	target=get_hexagon_target("v69"),
	)

	# Save a local copy of the Hexagon object code (in the form of a .so file)
	# to allow post-mortem inspection.
	host_dso_binary_path = os.path.join(host_files_dir_path, "test_binary.so")
	built_module.write_to_file(host_dso_binary_path)
	print(f"SAVED BINARY TO HOST PATH: {host_dso_binary_path}")

	hexagon_mod = hexagon_session.load_module(built_module)

	# Generate the input tensor's data.
	# Note that we'll eventually need it in two different layouts:
	# (1) NHWC as an argument to testing.poolnd_python.
	# (2) NHWC_8h8w32c for as an argument to our Hexagon primfunc.
	# a_numpy_4d = np.random.randint(low=-128, high=127,
	# size=input_tensor_shape_4d, dtype=np.int8)
	a_numpy_4d = _create_test_input(input_tensor_shape_4d, dtype)

	ref_output_4d = testing.poolnd_python(
	a_numpy_4d.astype("int32"),
	kernel,
	stride,
	dilation,
	padding[0:2],
	padding[2:],
	pool_type="max",
	dtype="int32",
	layout="NHWC",
	).astype(dtype)

	output_tensor_shape_4d = ref_output_4d.shape

	a_numpy_7d = _int8_nhwc_8h8w32c_xform_immediate(a_numpy_4d)

	a_hexagon_7d = allocate_hexagon_array(
	hexagon_session.device,
	tensor_shape=input_tensor_shape_7d,
	axis_separators=[4],
	dtype=dtype,
	mem_scope=io_tensor_mem_scope,
	)

	c_hexagon_4d = allocate_hexagon_array(
	hexagon_session.device,
	tensor_shape=output_tensor_shape_4d,
	axis_separators=[],
	dtype=dtype,
	mem_scope=io_tensor_mem_scope,
	)

	a_hexagon_7d.copyfrom(a_numpy_7d)

	if dtype == "int8":
	rel_tolerance = 0
	abs_tolerance = 0
	else:
	assert False, f"TODO: decide acceptable tolerances for dtype {dtype}"

	timer = hexagon_mod.time_evaluator(
	"main", hexagon_session.device, number=10, repeat=1
	)
	timing_result = timer(a_hexagon_7d, c_hexagon_4d)

	try:
	tvm.testing.assert_allclose(
	ref_output_4d, c_hexagon_4d.numpy(), rtol=rel_tolerance, atol=abs_tolerance
	)
	except AssertionError as exception:
	raise bu.NumericalAccuracyException(str(exception))

	except bu.NumericalAccuracyException as exception:
	print()
	print("FAIL: Numerical accuracy error. See log file.")

	log_file.write("\n")
	log_file.write(f"FAIL: {exception}\n")

	self.benchmark_table.record_fail(
	**keys_dict, comments="Numerical accuracy error. See log file."
	)

	except bu.UnsupportedException as exception:
	print()
	print(f"SKIP: {exception}")

	log_file.write("\n")
	log_file.write(f"SKIP: {exception}\n")

	self.benchmark_table.record_skip(
	**keys_dict, comments=f"Unsupported configuration: {exception}"
	)

	self.benchmark_table.record_success(timing_result, **keys_dict)


	if __name__ == "__main__":
	tvm.testing.main()