python/tvm/auto_scheduler/feature.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.

 """"
 Python API for Feature extraction. The extracted features vector are used by cost models.

 We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
 so we call this feature as "per-store" feature.
 The cost model also does prediction for each BufferStoreNode statement and aggregates
 the predicted score of each BufferStoreNode as the score of a TIR Stmt.

 The feature specification is defined by `src/auto_scheduler/feature.cc::FeatureSet`
 """

 from typing import List, Tuple, Union, Optional
 import struct

 import numpy as np

 from .loop_state import State, StateObject
 from .measure import MeasureInput, MeasureResult
 from . import _ffi_api

 # The maximum number of extracted buffers for one statement
 DEFAULT_MAX_N_BUFS = 5

 # The length of the feature vector
 DEFAULT_FEATURE_VEC_LEN = 164

 # The size of int and float in bytes
 SIZE_OF_INT32 = 4
 SIZE_OF_FLOAT32 = 4


 def unpack_feature(byte_arr: bytearray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Unpack the flatten feature (in byte array format) from c++

     Parameters
     ----------
     byte_arr: bytearray
         The two-dimensional feature vector in serialized byte array format

     Returns
     -------
     features: np.ndarray
         Feature vectors
     normalized_throughputs: np.ndarray
         Normalized throughputs
     task_ids: np.ndarray
         Task ids

     Note
     ----
     For faster data copy between c++ and python, the c++ part returns features in a single
     flatten array using a packed format. The python part then unpacks the flatten array.

     The packed format for n records is:
     {
       int   n;
       int   sizes[n+2];           // The sizes for the following arrays

       float features_0[size[0]];  // The features for record 0
       float features_1[size[1]];  // The features for record 1
       ...
       float features_i[size[i]];  // The features for record i
       ... // until i == n - 1

       float throughputs[sizes[n]];  // The normalized throughputs for n records
       int   task_ids[size[n+1];   // The task ids for n records

     }
     To implement this format, we also store int as float, so we can store all numbers
     into a single float array.
     """
     vec_len = DEFAULT_FEATURE_VEC_LEN

     # unpack sizes
     offset = 0
     n = struct.unpack_from("1i", byte_arr, offset=offset)[0]
     offset += SIZE_OF_INT32

     sizes = struct.unpack_from("%di" % (n + 2), byte_arr, offset=offset)
     offset += SIZE_OF_INT32 * (n + 2)

     # unpack features
     features = []
     for size in sizes[:-2]:
         row = []

         # Now, we need to unpack the feature for multiple statements.
         # The format is:
         # {
         #   int   n_stage;                        // The number of stages
         #   float feature_vecs[n_stage][vec_len]  // The feature vector for each stage
         # }
         # where vec_len can be calculated by `(size - 1) / n_stmts`

         if size == 0:
             # failed during lowering
             features.append(np.zeros((1, vec_len)))
         else:
             n_stmts = struct.unpack_from("f", byte_arr, offset=offset)
             offset += SIZE_OF_FLOAT32

             n_stmts = int(n_stmts[0] + 0.5)
             tmp_vec_len = (size - 1) // n_stmts
             assert (
                 tmp_vec_len == vec_len
             ), "The lenght of feature vector is wrong. " "Expected %d but got %d." % (
                 vec_len,
                 tmp_vec_len,
             )
             assert tmp_vec_len * n_stmts == size - 1
             for _ in range(n_stmts):
                 x = struct.unpack_from("%df" % vec_len, byte_arr, offset=offset)
                 offset += vec_len * SIZE_OF_FLOAT32
                 row.append(x)

             features.append(np.array(row))

     # unpack normalized_throughputs
     m = sizes[-2]
     normalized_throughputs = struct.unpack_from("%df" % m, byte_arr, offset=offset)
     offset += m * SIZE_OF_INT32

     # unpack task_ids
     m = sizes[-1]
     task_ids = struct.unpack_from("%di" % m, byte_arr, offset=offset)
     offset += m * SIZE_OF_INT32

     assert offset == len(byte_arr), "%d vs %d" % (offset, len(byte_arr))
     return np.array(features, dtype=object), np.array(normalized_throughputs), np.array(task_ids)


 def get_per_store_features_from_file(
     filename: str, max_lines: int, max_n_bufs: Optional[int] = None
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Get per-store features from a log file

     Parameters
     ----------
     filename: str
         The input filename
     max_lines: int
         Only extract the first n lines of the file
     max_n_bufs: Optional[int]
         The maximum number of extracted buffers for one statement

     Returns
     -------
     features: np.ndarray
         Feature vectors
     normalized_throughputs: np.ndarray
         Normalized throughputs
     task_ids: np.ndarray
         Task ids
     """
     byte_arr = _ffi_api.GetPerStoreFeaturesFromFile(
         filename, max_lines, max_n_bufs or DEFAULT_MAX_N_BUFS
     )
     return unpack_feature(byte_arr)


 def get_per_store_features_from_measure_pairs(
     inputs: List[MeasureInput],
     results: List[MeasureResult],
     skip_first_n_feature_extraction: int = 0,
     max_n_bufs: Optional[int] = None,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Get per-store features from measurement input/result pairs

     Parameters
     ----------
     inputs: List[MeasureInput]
         The measure inputs
     results: List[MeasureResult]
         The measure results
     skip_first_n_feature_extraction: int
         Skip feature extraction for the first n states
     max_n_bufs: int
         The maximum number of extracted buffers for one statement

     Returns
     -------
     features: np.ndarray
         Feature vectors
     normalized_throughputs: np.ndarray
         Normalized throughputs
     task_ids: np.ndarray
         Task ids
     """
     byte_arr = _ffi_api.GetPerStoreFeaturesFromMeasurePairs(
         inputs, results, skip_first_n_feature_extraction, max_n_bufs or DEFAULT_MAX_N_BUFS
     )
     return unpack_feature(byte_arr)


 def get_per_store_features_from_states(
     states: List[Union[State, StateObject]], task: "SearchTask", max_n_bufs: Optional[int] = None
 ) -> List[np.ndarray]:
     """Get per-store features from measurement input/result pairs

     Parameters
     ----------
     states: List[Union[State, StateObject]]
         The input states
     task: SearchTask
         The search task of the input states
     max_n_bufs: Optional[int]
         The maximum number of extracted buffers for one statement

     Returns
     -------
     features: np.ndarray
         Feature vectors
     normalized_throughputs: np.ndarray
         Normalized throughputs
     task_ids: np.ndarray
         Task ids
     """
     if isinstance(states[0], State):
         state_objects = [s.state_object for s in states]
     elif isinstance(states[0], StateObject):
         state_objects = states
     byte_arr = _ffi_api.GetPerStoreFeaturesFromStates(
         state_objects, task, max_n_bufs or DEFAULT_MAX_N_BUFS
     )
     return unpack_feature(byte_arr)[0]


 def get_per_store_feature_names(max_n_bufs: Optional[int] = None) -> List[str]:
     """Get the name of every element in the feature vector. Use this for debug and inspection.

     Parameters
     ----------
     max_n_bufs: int
         The maximum number of extracted buffers for one statement

     Returns
     -------
     names: List[str]
         The names of elements in the flatten feature vector
     """
     return _ffi_api.GetPerStoreFeatureNames(max_n_bufs or DEFAULT_MAX_N_BUFS)
	# 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.

	""""
	Python API for Feature extraction. The extracted features vector are used by cost models.

	We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
	so we call this feature as "per-store" feature.
	The cost model also does prediction for each BufferStoreNode statement and aggregates
	the predicted score of each BufferStoreNode as the score of a TIR Stmt.

	The feature specification is defined by `src/auto_scheduler/feature.cc::FeatureSet`
	"""

	from typing import List, Tuple, Union, Optional
	import struct

	import numpy as np

	from .loop_state import State, StateObject
	from .measure import MeasureInput, MeasureResult
	from . import _ffi_api

	# The maximum number of extracted buffers for one statement
	DEFAULT_MAX_N_BUFS = 5

	# The length of the feature vector
	DEFAULT_FEATURE_VEC_LEN = 164

	# The size of int and float in bytes
	SIZE_OF_INT32 = 4
	SIZE_OF_FLOAT32 = 4


	def unpack_feature(byte_arr: bytearray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
	"""Unpack the flatten feature (in byte array format) from c++

	Parameters
	----------
	byte_arr: bytearray
	The two-dimensional feature vector in serialized byte array format

	Returns
	-------
	features: np.ndarray
	Feature vectors
	normalized_throughputs: np.ndarray
	Normalized throughputs
	task_ids: np.ndarray
	Task ids

	Note
	----
	For faster data copy between c++ and python, the c++ part returns features in a single
	flatten array using a packed format. The python part then unpacks the flatten array.

	The packed format for n records is:
	{
	int n;
	int sizes[n+2]; // The sizes for the following arrays

	float features_0[size[0]]; // The features for record 0
	float features_1[size[1]]; // The features for record 1
	...
	float features_i[size[i]]; // The features for record i
	... // until i == n - 1

	float throughputs[sizes[n]]; // The normalized throughputs for n records
	int task_ids[size[n+1]; // The task ids for n records

	}
	To implement this format, we also store int as float, so we can store all numbers
	into a single float array.
	"""
	vec_len = DEFAULT_FEATURE_VEC_LEN

	# unpack sizes
	offset = 0
	n = struct.unpack_from("1i", byte_arr, offset=offset)[0]
	offset += SIZE_OF_INT32

	sizes = struct.unpack_from("%di" % (n + 2), byte_arr, offset=offset)
	offset += SIZE_OF_INT32 * (n + 2)

	# unpack features
	features = []
	for size in sizes[:-2]:
	row = []

	# Now, we need to unpack the feature for multiple statements.
	# The format is:
	# {
	# int n_stage; // The number of stages
	# float feature_vecs[n_stage][vec_len] // The feature vector for each stage
	# }
	# where vec_len can be calculated by `(size - 1) / n_stmts`

	if size == 0:
	# failed during lowering
	features.append(np.zeros((1, vec_len)))
	else:
	n_stmts = struct.unpack_from("f", byte_arr, offset=offset)
	offset += SIZE_OF_FLOAT32

	n_stmts = int(n_stmts[0] + 0.5)
	tmp_vec_len = (size - 1) // n_stmts
	assert (
	tmp_vec_len == vec_len
	), "The lenght of feature vector is wrong. " "Expected %d but got %d." % (
	vec_len,
	tmp_vec_len,
	)
	assert tmp_vec_len * n_stmts == size - 1
	for _ in range(n_stmts):
	x = struct.unpack_from("%df" % vec_len, byte_arr, offset=offset)
	offset += vec_len * SIZE_OF_FLOAT32
	row.append(x)

	features.append(np.array(row))

	# unpack normalized_throughputs
	m = sizes[-2]
	normalized_throughputs = struct.unpack_from("%df" % m, byte_arr, offset=offset)
	offset += m * SIZE_OF_INT32

	# unpack task_ids
	m = sizes[-1]
	task_ids = struct.unpack_from("%di" % m, byte_arr, offset=offset)
	offset += m * SIZE_OF_INT32

	assert offset == len(byte_arr), "%d vs %d" % (offset, len(byte_arr))
	return np.array(features, dtype=object), np.array(normalized_throughputs), np.array(task_ids)


	def get_per_store_features_from_file(
	filename: str, max_lines: int, max_n_bufs: Optional[int] = None
	) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
	"""Get per-store features from a log file

	Parameters
	----------
	filename: str
	The input filename
	max_lines: int
	Only extract the first n lines of the file
	max_n_bufs: Optional[int]
	The maximum number of extracted buffers for one statement

	Returns
	-------
	features: np.ndarray
	Feature vectors
	normalized_throughputs: np.ndarray
	Normalized throughputs
	task_ids: np.ndarray
	Task ids
	"""
	byte_arr = _ffi_api.GetPerStoreFeaturesFromFile(
	filename, max_lines, max_n_bufs or DEFAULT_MAX_N_BUFS
	)
	return unpack_feature(byte_arr)


	def get_per_store_features_from_measure_pairs(
	inputs: List[MeasureInput],
	results: List[MeasureResult],
	skip_first_n_feature_extraction: int = 0,
	max_n_bufs: Optional[int] = None,
	) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
	"""Get per-store features from measurement input/result pairs

	Parameters
	----------
	inputs: List[MeasureInput]
	The measure inputs
	results: List[MeasureResult]
	The measure results
	skip_first_n_feature_extraction: int
	Skip feature extraction for the first n states
	max_n_bufs: int
	The maximum number of extracted buffers for one statement

	Returns
	-------
	features: np.ndarray
	Feature vectors
	normalized_throughputs: np.ndarray
	Normalized throughputs
	task_ids: np.ndarray
	Task ids
	"""
	byte_arr = _ffi_api.GetPerStoreFeaturesFromMeasurePairs(
	inputs, results, skip_first_n_feature_extraction, max_n_bufs or DEFAULT_MAX_N_BUFS
	)
	return unpack_feature(byte_arr)


	def get_per_store_features_from_states(
	states: List[Union[State, StateObject]], task: "SearchTask", max_n_bufs: Optional[int] = None
	) -> List[np.ndarray]:
	"""Get per-store features from measurement input/result pairs

	Parameters
	----------
	states: List[Union[State, StateObject]]
	The input states
	task: SearchTask
	The search task of the input states
	max_n_bufs: Optional[int]
	The maximum number of extracted buffers for one statement

	Returns
	-------
	features: np.ndarray
	Feature vectors
	normalized_throughputs: np.ndarray
	Normalized throughputs
	task_ids: np.ndarray
	Task ids
	"""
	if isinstance(states[0], State):
	state_objects = [s.state_object for s in states]
	elif isinstance(states[0], StateObject):
	state_objects = states
	byte_arr = _ffi_api.GetPerStoreFeaturesFromStates(
	state_objects, task, max_n_bufs or DEFAULT_MAX_N_BUFS
	)
	return unpack_feature(byte_arr)[0]


	def get_per_store_feature_names(max_n_bufs: Optional[int] = None) -> List[str]:
	"""Get the name of every element in the feature vector. Use this for debug and inspection.

	Parameters
	----------
	max_n_bufs: int
	The maximum number of extracted buffers for one statement

	Returns
	-------
	names: List[str]
	The names of elements in the flatten feature vector
	"""
	return _ffi_api.GetPerStoreFeatureNames(max_n_bufs or DEFAULT_MAX_N_BUFS)