chore: import upstream snapshot with attribution

python/tvm/s_tir/meta_schedule/feature_extractor/__init__.py

@@ -0,0 +1,26 @@
+# isort: skip_file
+# 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.
+"""
+The tvm.s_tir.meta_schedule.feature_extractor package.
+Meta Schedule feature extractors that extracts features from
+measure candidates for use in cost model.
+"""
+
+from .feature_extractor import FeatureExtractor, PyFeatureExtractor
+from .per_store_feature import PerStoreFeature
+from .random_feature_extractor import RandomFeatureExtractor

python/tvm/s_tir/meta_schedule/feature_extractor/feature_extractor.py

@@ -0,0 +1,128 @@
+# 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: RUF012
+"""Meta Schedule FeatureExtractor."""
+
+from collections.abc import Callable
+from typing import Union
+
+# isort: off
+from typing import Literal
+
+# isort: on
+
+from tvm_ffi import register_object
+
+from tvm.runtime import Object
+from tvm.runtime._tensor import Tensor
+
+from .. import _ffi_api
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@register_object("s_tir.meta_schedule.FeatureExtractor")
+class FeatureExtractor(Object):
+    """Extractor for features from measure candidates for use in cost model."""
+
+    FeatureExtractorType = Union[Literal["per-store-feature"], "FeatureExtractor"]
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[Tensor]:
+        """Extract features from the given measure candidate.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for feature extraction.
+        candidates : List[MeasureCandidate]
+            The measure candidates to extract features from.
+
+        Returns
+        -------
+        features : List[Tensor]
+            The feature tvm ndarray extracted.
+        """
+        result = _ffi_api.FeatureExtractorExtractFrom(  # type: ignore # pylint: disable=no-member
+            self, context, candidates
+        )
+        return result
+
+    @staticmethod
+    def create(
+        kind: Literal["per-store-feature"],
+        *args,
+        **kwargs,
+    ) -> "FeatureExtractor":
+        """Create a CostModel."""
+        from . import PerStoreFeature  # pylint: disable=import-outside-toplevel
+
+        if kind == "per-store-feature":
+            return PerStoreFeature(*args, **kwargs)  # type: ignore
+        raise ValueError(f"Unknown CostModel: {kind}")
+
+
+@register_object("s_tir.meta_schedule.PyFeatureExtractor")
+class _PyFeatureExtractor(FeatureExtractor):
+    """
+    A TVM object feature extractor to support customization on the python side.
+    This is NOT the user facing class for function overloading inheritance.
+
+    See also: PyFeatureExtractor
+    """
+
+    def __init__(self, f_extract_from: Callable):
+        """Constructor."""
+
+        self.__init_handle_by_constructor__(
+            _ffi_api.FeatureExtractorPyFeatureExtractor,  # type: ignore # pylint: disable=no-member
+            f_extract_from,
+        )
+
+
+class PyFeatureExtractor:
+    """
+    An abstract feature extractor with customized methods on the python-side.
+    This is the user facing class for function overloading inheritance.
+
+    Note: @derived_object is required for proper usage of any inherited class.
+    """
+
+    _tvm_metadata = {
+        "cls": _PyFeatureExtractor,
+        "methods": ["extract_from"],
+    }
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[Tensor]:
+        """Extract features from the given measure candidate.
+
+        Parameters
+        ----------
+        context : TuneContext
+            The tuning context for feature extraction.
+        candidates : List[MeasureCandidate]
+            The measure candidates to extract features from.
+
+        Returns
+        -------
+        features : List[Tensor]
+            The feature tvm ndarray extracted.
+        """
+        raise NotImplementedError

python/tvm/s_tir/meta_schedule/feature_extractor/per_store_feature.py

@@ -0,0 +1,68 @@
+# 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.
+# pylint: disable=invalid-name
+"""We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
+so we call this feature as "per-store" feature.
+"""
+
+from tvm_ffi import register_object
+
+from .. import _ffi_api
+from .feature_extractor import FeatureExtractor
+
+
+@register_object("s_tir.meta_schedule.PerStoreFeature")
+class PerStoreFeature(FeatureExtractor):
+    """PerStoreFeature extracts one feature vector per BufferStoreNode
+
+    Parameters
+    ----------
+    buffers_per_store : int
+        The number of buffers in each BufferStore; Pad or truncate if necessary.
+    arith_intensity_curve_num_samples : int
+        The number of samples used in the arithmetic intensity curve.
+    cache_line_bytes : int
+        The number of bytes in a cache line.
+    extract_workload : bool
+        Whether to extract features in the workload in tuning context or not.
+    """
+
+    buffers_per_store: int
+    """The number of buffers in each BufferStore; Pad or truncate if necessary."""
+    arith_intensity_curve_num_samples: int  # pylint: disable=invalid-name
+    """The number of samples used in the arithmetic intensity curve."""
+    cache_line_bytes: int
+    """The number of bytes in a cache line."""
+    extract_workload: bool
+    """Whether to extract features in the workload in tuning context or not."""
+    feature_vector_length: int
+    """Length of the feature vector."""
+
+    def __init__(
+        self,
+        buffers_per_store: int = 5,
+        arith_intensity_curve_num_samples: int = 10,
+        cache_line_bytes: int = 64,
+        extract_workload: bool = False,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.FeatureExtractorPerStoreFeature,  # type: ignore # pylint: disable=no-member
+            buffers_per_store,
+            arith_intensity_curve_num_samples,
+            cache_line_bytes,
+            extract_workload,
+        )

python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py

@@ -0,0 +1,64 @@
+# 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.
+"""Random Feature Extractor."""
+
+import numpy as np  # type: ignore
+
+import tvm.runtime
+from tvm.ir.utils import derived_object
+
+from ..feature_extractor import PyFeatureExtractor
+from ..search_strategy import MeasureCandidate
+from ..tune_context import TuneContext
+
+
+@derived_object
+class RandomFeatureExtractor(PyFeatureExtractor):
+    """Random Feature Extractor
+
+    Parameters
+    ----------
+    feature_size : int
+        The size of each block's feature vector.
+    max_block_num : int
+        The maximum number of blocks in each schedule.
+    random_state : Union[Tuple[str, np.ndarray, int, int, float], dict]
+        The current random state of the f
+    """
+
+    feature_size: int
+    max_block_num: int
+    random_state: tuple[str, np.ndarray, int, int, float] | dict
+
+    def __init__(self, *, feature_size: int = 30, max_block_num: int = 5, seed=0):
+        super().__init__()
+        assert max_block_num >= 1, "Max block number must be greater or equal to one!"
+        self.max_block_num = max_block_num
+        self.feature_size = feature_size
+        np.random.seed(seed)
+        self.random_state = np.random.get_state()
+
+    def extract_from(
+        self, context: TuneContext, candidates: list[MeasureCandidate]
+    ) -> list[tvm.runtime.Tensor]:
+        np.random.set_state(self.random_state)
+        result = [
+            np.random.rand(np.random.randint(1, self.max_block_num + 1), self.feature_size)
+            for candidate in candidates
+        ]
+        self.random_state = np.random.get_state()
+        return [tvm.runtime.tensor(x) for x in result]