devtools/bundled_program/util/test_util.py - platform/external/executorch - Git at Google

 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.

 # pyre-strict
 import random
 import string
 from typing import List, Tuple

 import torch
 from executorch.devtools.bundled_program.config import (
     MethodInputType,
     MethodOutputType,
     MethodTestCase,
     MethodTestSuite,
 )

 from executorch.exir import ExecutorchProgramManager, to_edge
 from torch.export import export
 from torch.export.unflatten import _assign_attr, _AttrKind

 # A hacky integer to deal with a mismatch between execution plan and complier.
 #
 # Execution plan supports multiple types of inputs, like Tensor, Int, etc,
 # rather than only Tensor. However, compiler only supports Tensor as input type.
 # All other inputs will remain the same as default value in the model, which
 # means during model execution, each function will use the preset default value
 # for non-tensor inputs, rather than the one we manually set. However, eager
 # model supports multiple types of inputs.
 #
 # In order to show that bundled program can support multiple input types while
 # executorch model can generate the same output as eager model, we hackily set
 # all Int inputs in Bundled Program and default int inputs in model as a same
 # value, called DEFAULT_INT_INPUT.
 #
 # TODO(gasoonjia): track the situation. Stop supporting multiple input types in
 # bundled program if execution plan stops supporting it, or remove this hacky
 # method if compiler can support multiple input types
 DEFAULT_INT_INPUT = 2


 class SampleModel(torch.nn.Module):
     """An example model with multi-methods. Each method has multiple input and single output"""

     def __init__(self) -> None:
         super().__init__()
         self.register_buffer("a", 3 * torch.ones(2, 2, dtype=torch.int32))
         self.register_buffer("b", 2 * torch.ones(2, 2, dtype=torch.int32))
         self.method_names = ["encode", "decode"]

     def encode(
         self, x: torch.Tensor, q: torch.Tensor, a: int = DEFAULT_INT_INPUT
     ) -> torch.Tensor:
         z = x.clone()
         torch.mul(self.a, x, out=z)
         y = x.clone()
         torch.add(z, self.b, alpha=a, out=y)
         torch.add(y, q, out=y)
         return y

     def decode(
         self, x: torch.Tensor, q: torch.Tensor, a: int = DEFAULT_INT_INPUT
     ) -> torch.Tensor:
         y = x * q
         torch.add(y, self.b, alpha=a, out=y)
         return y


 def get_rand_input_values(
     n_tensors: int,
     sizes: List[List[int]],
     n_int: int,
     dtype: torch.dtype,
     n_sets_per_plan_test: int,
     n_method_test_suites: int,
 ) -> List[List[MethodInputType]]:
     # pyre-ignore[7]: expected `List[List[List[Union[bool, float, int, Tensor]]]]` but got `List[List[List[Union[int, Tensor]]]]`
     return [
         [
             [(torch.rand(*sizes[i]) - 0.5).to(dtype) for i in range(n_tensors)]
             + [DEFAULT_INT_INPUT for _ in range(n_int)]
             for _ in range(n_sets_per_plan_test)
         ]
         for _ in range(n_method_test_suites)
     ]


 def get_rand_output_values(
     n_tensors: int,
     sizes: List[List[int]],
     dtype: torch.dtype,
     n_sets_per_plan_test: int,
     n_method_test_suites: int,
 ) -> List[List[MethodOutputType]]:
     # pyre-ignore [7]: Expected `List[List[Sequence[Tensor]]]` but got `List[List[List[Tensor]]]`.
     return [
         [
             [(torch.rand(*sizes[i]) - 0.5).to(dtype) for i in range(n_tensors)]
             for _ in range(n_sets_per_plan_test)
         ]
         for _ in range(n_method_test_suites)
     ]


 def get_rand_method_names(n_method_test_suites: int) -> List[str]:
     unique_strings = set()
     while len(unique_strings) < n_method_test_suites:
         rand_str = "".join(random.choices(string.ascii_letters, k=5))
         if rand_str not in unique_strings:
             unique_strings.add(rand_str)
     return list(unique_strings)


 def get_random_test_suites(
     n_model_inputs: int,
     model_input_sizes: List[List[int]],
     n_model_outputs: int,
     model_output_sizes: List[List[int]],
     dtype: torch.dtype,
     n_sets_per_plan_test: int,
     n_method_test_suites: int,
 ) -> Tuple[
     List[str],
     List[List[MethodInputType]],
     List[List[MethodOutputType]],
     List[MethodTestSuite],
 ]:
     """Helper function to generate config filled with random inputs and expected outputs.

     The return type of rand inputs is a List[List[InputValues]]. The inner list of
     InputValues represents all test sets for single execution plan, while the outer list
     is for multiple execution plans.

     Same for rand_expected_outputs.

     """

     rand_method_names = get_rand_method_names(n_method_test_suites)

     rand_inputs_per_program = get_rand_input_values(
         n_tensors=n_model_inputs,
         sizes=model_input_sizes,
         n_int=1,
         dtype=dtype,
         n_sets_per_plan_test=n_sets_per_plan_test,
         n_method_test_suites=n_method_test_suites,
     )

     rand_expected_output_per_program = get_rand_output_values(
         n_tensors=n_model_outputs,
         sizes=model_output_sizes,
         dtype=dtype,
         n_sets_per_plan_test=n_sets_per_plan_test,
         n_method_test_suites=n_method_test_suites,
     )

     rand_method_test_suites: List[MethodTestSuite] = []

     for (
         rand_method_name,
         rand_inputs_per_method,
         rand_expected_output_per_method,
     ) in zip(
         rand_method_names, rand_inputs_per_program, rand_expected_output_per_program
     ):
         rand_method_test_cases: List[MethodTestCase] = []
         for rand_inputs, rand_expected_outputs in zip(
             rand_inputs_per_method, rand_expected_output_per_method
         ):
             rand_method_test_cases.append(
                 MethodTestCase(
                     inputs=rand_inputs, expected_outputs=rand_expected_outputs
                 )
             )

         rand_method_test_suites.append(
             MethodTestSuite(
                 method_name=rand_method_name, test_cases=rand_method_test_cases
             )
         )

     return (
         rand_method_names,
         rand_inputs_per_program,
         rand_expected_output_per_program,
         rand_method_test_suites,
     )


 def get_random_test_suites_with_eager_model(
     eager_model: torch.nn.Module,
     method_names: List[str],
     n_model_inputs: int,
     model_input_sizes: List[List[int]],
     dtype: torch.dtype,
     n_sets_per_plan_test: int,
 ) -> Tuple[List[List[MethodInputType]], List[MethodTestSuite]]:
     """Generate config filled with random inputs for each inference method given eager model

     The details of return type is the same as get_random_test_suites_with_rand_io_lists.
     """
     inputs_per_program = get_rand_input_values(
         n_tensors=n_model_inputs,
         sizes=model_input_sizes,
         n_int=1,
         dtype=dtype,
         n_sets_per_plan_test=n_sets_per_plan_test,
         n_method_test_suites=len(method_names),
     )

     method_test_suites: List[MethodTestSuite] = []

     for method_name, inputs_per_method in zip(method_names, inputs_per_program):
         method_test_cases: List[MethodTestCase] = []
         for inputs in inputs_per_method:
             method_test_cases.append(
                 MethodTestCase(
                     inputs=inputs,
                     expected_outputs=getattr(eager_model, method_name)(*inputs),
                 )
             )

         method_test_suites.append(
             MethodTestSuite(method_name=method_name, test_cases=method_test_cases)
         )

     return inputs_per_program, method_test_suites


 class StatefulWrapperModule(torch.nn.Module):
     """A version of wrapper module that preserves parameters/buffers.

     Use this if you are planning to wrap a non-forward method on an existing
     module.
     """

     def __init__(self, base_mod, method) -> None:  # pyre-ignore
         super().__init__()
         state_dict = base_mod.state_dict()
         for name, value in base_mod.named_parameters():
             _assign_attr(value, self, name, _AttrKind.PARAMETER)
         for name, value in base_mod.named_buffers():
             _assign_attr(
                 value, self, name, _AttrKind.BUFFER, persistent=name in state_dict
             )
         self.fn = method  # pyre-ignore

     def forward(self, *args, **kwargs):  # pyre-ignore
         return self.fn(*args, **kwargs)


 def get_common_executorch_program() -> (
     Tuple[ExecutorchProgramManager, List[MethodTestSuite]]
 ):
     """Helper function to generate a sample BundledProgram with its config."""
     eager_model = SampleModel()
     # Trace to FX Graph.
     capture_inputs = {
         m_name: (
             (torch.rand(2, 2) - 0.5).to(dtype=torch.int32),
             (torch.rand(2, 2) - 0.5).to(dtype=torch.int32),
             DEFAULT_INT_INPUT,
         )
         for m_name in eager_model.method_names
     }

     # Trace to FX Graph and emit the program
     method_graphs = {
         m_name: export(
             StatefulWrapperModule(eager_model, getattr(eager_model, m_name)),
             capture_inputs[m_name],
         )
         for m_name in eager_model.method_names
     }

     executorch_program = to_edge(method_graphs).to_executorch()

     _, method_test_suites = get_random_test_suites_with_eager_model(
         eager_model=eager_model,
         method_names=eager_model.method_names,
         n_model_inputs=2,
         model_input_sizes=[[2, 2], [2, 2]],
         dtype=torch.int32,
         n_sets_per_plan_test=10,
     )
     return executorch_program, method_test_suites
	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.

	# pyre-strict
	import random
	import string
	from typing import List, Tuple

	import torch
	from executorch.devtools.bundled_program.config import (
	MethodInputType,
	MethodOutputType,
	MethodTestCase,
	MethodTestSuite,
	)

	from executorch.exir import ExecutorchProgramManager, to_edge
	from torch.export import export
	from torch.export.unflatten import _assign_attr, _AttrKind

	# A hacky integer to deal with a mismatch between execution plan and complier.
	#
	# Execution plan supports multiple types of inputs, like Tensor, Int, etc,
	# rather than only Tensor. However, compiler only supports Tensor as input type.
	# All other inputs will remain the same as default value in the model, which
	# means during model execution, each function will use the preset default value
	# for non-tensor inputs, rather than the one we manually set. However, eager
	# model supports multiple types of inputs.
	#
	# In order to show that bundled program can support multiple input types while
	# executorch model can generate the same output as eager model, we hackily set
	# all Int inputs in Bundled Program and default int inputs in model as a same
	# value, called DEFAULT_INT_INPUT.
	#
	# TODO(gasoonjia): track the situation. Stop supporting multiple input types in
	# bundled program if execution plan stops supporting it, or remove this hacky
	# method if compiler can support multiple input types
	DEFAULT_INT_INPUT = 2


	class SampleModel(torch.nn.Module):
	"""An example model with multi-methods. Each method has multiple input and single output"""

	def __init__(self) -> None:
	super().__init__()
	self.register_buffer("a", 3 * torch.ones(2, 2, dtype=torch.int32))
	self.register_buffer("b", 2 * torch.ones(2, 2, dtype=torch.int32))
	self.method_names = ["encode", "decode"]

	def encode(
	self, x: torch.Tensor, q: torch.Tensor, a: int = DEFAULT_INT_INPUT
	) -> torch.Tensor:
	z = x.clone()
	torch.mul(self.a, x, out=z)
	y = x.clone()
	torch.add(z, self.b, alpha=a, out=y)
	torch.add(y, q, out=y)
	return y

	def decode(
	self, x: torch.Tensor, q: torch.Tensor, a: int = DEFAULT_INT_INPUT
	) -> torch.Tensor:
	y = x * q
	torch.add(y, self.b, alpha=a, out=y)
	return y


	def get_rand_input_values(
	n_tensors: int,
	sizes: List[List[int]],
	n_int: int,
	dtype: torch.dtype,
	n_sets_per_plan_test: int,
	n_method_test_suites: int,
	) -> List[List[MethodInputType]]:
	# pyre-ignore[7]: expected `List[List[List[Union[bool, float, int, Tensor]]]]` but got `List[List[List[Union[int, Tensor]]]]`
	return [
	[
	[(torch.rand(*sizes[i]) - 0.5).to(dtype) for i in range(n_tensors)]
	+ [DEFAULT_INT_INPUT for _ in range(n_int)]
	for _ in range(n_sets_per_plan_test)
	]
	for _ in range(n_method_test_suites)
	]


	def get_rand_output_values(
	n_tensors: int,
	sizes: List[List[int]],
	dtype: torch.dtype,
	n_sets_per_plan_test: int,
	n_method_test_suites: int,
	) -> List[List[MethodOutputType]]:
	# pyre-ignore [7]: Expected `List[List[Sequence[Tensor]]]` but got `List[List[List[Tensor]]]`.
	return [
	[
	[(torch.rand(*sizes[i]) - 0.5).to(dtype) for i in range(n_tensors)]
	for _ in range(n_sets_per_plan_test)
	]
	for _ in range(n_method_test_suites)
	]


	def get_rand_method_names(n_method_test_suites: int) -> List[str]:
	unique_strings = set()
	while len(unique_strings) < n_method_test_suites:
	rand_str = "".join(random.choices(string.ascii_letters, k=5))
	if rand_str not in unique_strings:
	unique_strings.add(rand_str)
	return list(unique_strings)


	def get_random_test_suites(
	n_model_inputs: int,
	model_input_sizes: List[List[int]],
	n_model_outputs: int,
	model_output_sizes: List[List[int]],
	dtype: torch.dtype,
	n_sets_per_plan_test: int,
	n_method_test_suites: int,
	) -> Tuple[
	List[str],
	List[List[MethodInputType]],
	List[List[MethodOutputType]],
	List[MethodTestSuite],
	]:
	"""Helper function to generate config filled with random inputs and expected outputs.

	The return type of rand inputs is a List[List[InputValues]]. The inner list of
	InputValues represents all test sets for single execution plan, while the outer list
	is for multiple execution plans.

	Same for rand_expected_outputs.

	"""

	rand_method_names = get_rand_method_names(n_method_test_suites)

	rand_inputs_per_program = get_rand_input_values(
	n_tensors=n_model_inputs,
	sizes=model_input_sizes,
	n_int=1,
	dtype=dtype,
	n_sets_per_plan_test=n_sets_per_plan_test,
	n_method_test_suites=n_method_test_suites,
	)

	rand_expected_output_per_program = get_rand_output_values(
	n_tensors=n_model_outputs,
	sizes=model_output_sizes,
	dtype=dtype,
	n_sets_per_plan_test=n_sets_per_plan_test,
	n_method_test_suites=n_method_test_suites,
	)

	rand_method_test_suites: List[MethodTestSuite] = []

	for (
	rand_method_name,
	rand_inputs_per_method,
	rand_expected_output_per_method,
	) in zip(
	rand_method_names, rand_inputs_per_program, rand_expected_output_per_program
	):
	rand_method_test_cases: List[MethodTestCase] = []
	for rand_inputs, rand_expected_outputs in zip(
	rand_inputs_per_method, rand_expected_output_per_method
	):
	rand_method_test_cases.append(
	MethodTestCase(
	inputs=rand_inputs, expected_outputs=rand_expected_outputs
	)
	)

	rand_method_test_suites.append(
	MethodTestSuite(
	method_name=rand_method_name, test_cases=rand_method_test_cases
	)
	)

	return (
	rand_method_names,
	rand_inputs_per_program,
	rand_expected_output_per_program,
	rand_method_test_suites,
	)


	def get_random_test_suites_with_eager_model(
	eager_model: torch.nn.Module,
	method_names: List[str],
	n_model_inputs: int,
	model_input_sizes: List[List[int]],
	dtype: torch.dtype,
	n_sets_per_plan_test: int,
	) -> Tuple[List[List[MethodInputType]], List[MethodTestSuite]]:
	"""Generate config filled with random inputs for each inference method given eager model

	The details of return type is the same as get_random_test_suites_with_rand_io_lists.
	"""
	inputs_per_program = get_rand_input_values(
	n_tensors=n_model_inputs,
	sizes=model_input_sizes,
	n_int=1,
	dtype=dtype,
	n_sets_per_plan_test=n_sets_per_plan_test,
	n_method_test_suites=len(method_names),
	)

	method_test_suites: List[MethodTestSuite] = []

	for method_name, inputs_per_method in zip(method_names, inputs_per_program):
	method_test_cases: List[MethodTestCase] = []
	for inputs in inputs_per_method:
	method_test_cases.append(
	MethodTestCase(
	inputs=inputs,
	expected_outputs=getattr(eager_model, method_name)(*inputs),
	)
	)

	method_test_suites.append(
	MethodTestSuite(method_name=method_name, test_cases=method_test_cases)
	)

	return inputs_per_program, method_test_suites


	class StatefulWrapperModule(torch.nn.Module):
	"""A version of wrapper module that preserves parameters/buffers.

	Use this if you are planning to wrap a non-forward method on an existing
	module.
	"""

	def __init__(self, base_mod, method) -> None: # pyre-ignore
	super().__init__()
	state_dict = base_mod.state_dict()
	for name, value in base_mod.named_parameters():
	_assign_attr(value, self, name, _AttrKind.PARAMETER)
	for name, value in base_mod.named_buffers():
	_assign_attr(
	value, self, name, _AttrKind.BUFFER, persistent=name in state_dict
	)
	self.fn = method # pyre-ignore

	def forward(self, args, *kwargs): # pyre-ignore
	return self.fn(args, *kwargs)


	def get_common_executorch_program() -> (
	Tuple[ExecutorchProgramManager, List[MethodTestSuite]]
	):
	"""Helper function to generate a sample BundledProgram with its config."""
	eager_model = SampleModel()
	# Trace to FX Graph.
	capture_inputs = {
	m_name: (
	(torch.rand(2, 2) - 0.5).to(dtype=torch.int32),
	(torch.rand(2, 2) - 0.5).to(dtype=torch.int32),
	DEFAULT_INT_INPUT,
	)
	for m_name in eager_model.method_names
	}

	# Trace to FX Graph and emit the program
	method_graphs = {
	m_name: export(
	StatefulWrapperModule(eager_model, getattr(eager_model, m_name)),
	capture_inputs[m_name],
	)
	for m_name in eager_model.method_names
	}

	executorch_program = to_edge(method_graphs).to_executorch()

	_, method_test_suites = get_random_test_suites_with_eager_model(
	eager_model=eager_model,
	method_names=eager_model.method_names,
	n_model_inputs=2,
	model_input_sizes=[[2, 2], [2, 2]],
	dtype=torch.int32,
	n_sets_per_plan_test=10,
	)
	return executorch_program, method_test_suites