tkonolige commented on a change in pull request #10455:
URL: https://github.com/apache/tvm/pull/10455#discussion_r822959339
##########
File path: python/tvm/auto_scheduler/feature.py
##########
@@ -252,3 +253,78 @@ def get_per_store_feature_names(max_n_bufs: Optional[int]
= None) -> List[str]:
The names of elements in the flatten feature vector
"""
return _ffi_api.GetPerStoreFeatureNames(max_n_bufs or DEFAULT_MAX_N_BUFS)
+
+
+def primfunc_features(
Review comment:
Renamed to `features_from_primfunc`
##########
File path: python/tvm/auto_scheduler/feature.py
##########
@@ -252,3 +253,78 @@ def get_per_store_feature_names(max_n_bufs: Optional[int]
= None) -> List[str]:
The names of elements in the flatten feature vector
"""
return _ffi_api.GetPerStoreFeatureNames(max_n_bufs or DEFAULT_MAX_N_BUFS)
+
+
+def primfunc_features(
+ func: PrimFunc,
+ cache_line_bytes: int = 64,
+ max_n_bufs: Optional[int] = None,
+ log_scale: bool = False,
+) -> np.ndarray:
+ """Extract performance features from a PrimFunc.
+
+ Parameters
+ ----------
+ func: PrimFunc
+ PrimFunc from which features will be extracted. Each store operation to
+ a unique buffer in the function will result in one row of features in
+ the output.
+
+ cache_line_bytes: int, optional
+ Size of a cache line in bytes. Defaults to 64 which is the size for
+ most x86 processors.
+
+ max_n_bufs: int, optional
+ Maximum number of buffers in generated features. This determines the
+ length of the resulting feature vector.
+
+ log_scale: bool
+ Should entries in the feature vector be scaled by log2(x + 1). Defaults
+ to False. Use True if using features with a cost model.
+
+ Returns
+ -------
+ np.ndarray
+ Output features, one row per store into a unique buffer statement in
`func`.
+ """
+ return _ffi_api.FeaturesFromPrimFunc(
+ func, cache_line_bytes, max_n_bufs or DEFAULT_MAX_N_BUFS, log_scale
+ ).numpy()
+
+
+def named_primfunc_features(
Review comment:
Renamed to `named_features_from_primfunc`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
