# 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. """Data layout.""" import tvm_ffi from tvm.runtime import Object from . import _ffi_api @tvm_ffi.register_object("s_tir.SLayout") class SLayout(Object): """SLayout is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. For example, NCHW16c can describe a 5-D tensor of [batch_size, channel, height, width, channel_block]. Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel). See Also -------- slayout : Declare a layout """ def __len__(self): return _ffi_api.SLayoutNdim(self) # type: ignore def __contains__(self, axis): # Note: We do a weaker check for packed axis assuming layout is valid return not any(bkt in axis for bkt in "[]") and axis in self.name def __getitem__(self, index): if index >= len(self): raise IndexError("SLayout index out of range") return _ffi_api.SLayoutGetItem(self, index) # type: ignore def index_of(self, axis): """Get the index of an axis Parameters ---------- axis : str The axis name, needs to be [a-z,A-Z] or a packed axis Returns ------- index : int The index of the axis, -1 if not found. """ return _ffi_api.SLayoutIndexOf(self, axis) # type: ignore def factor_of(self, axis): """Get the factor size of the subordinate axis. Parameters ---------- axis : str The axis name, need to be [a-z,A-Z] Returns ------- factor : int the size of the subordinate-axis of axis (if axis is a primal-axis), or the size of axis itself (if axis is a subordinate-axis). Return -1 if axis is not in the layout. """ return _ffi_api.SLayoutFactorOf(self, axis) # type: ignore @tvm_ffi.register_object("s_tir.SBijectiveLayout") class SBijectiveLayout(Object): """Bijective mapping for two layouts (src-layout and dst-layout). It provides shape and index conversion between each other. Do not construct directly, use :any:`sbijective_layout` instead. See the documentation of :any:`sbijective_layout` for more details. Parameters ---------- src_layout : str or SLayout source layout. dst_layout : str or SLayout destination layout. See Also -------- sbijective_layout : Declare a layout """ def forward_index(self, index): """Given the indices of the src-layout, infer the dst index. Parameters ---------- index: Array of Expr The indices in src-layout. Returns ------- dst_index: Array of Expr The inferred indices in dst-layout. """ return _ffi_api.SBijectiveLayoutForwardIndex(self, index) # type: ignore def backward_index(self, index): """Given the indices of the dst-layout, infer the src index. Parameters ---------- index: Array of Expr The indices in dst-layout. Returns ------- src_index: Array of Expr The inferred indices in src-layout. """ return _ffi_api.SBijectiveLayoutBackwardIndex(self, index) # type: ignore def forward_shape(self, shape): """Given the shape of the src-layout, infer the dst shape. Parameters ---------- shape: Array of Expr The shape in src-layout. Returns ------- dst_shape: Array of Expr The inferred shape in dst-layout. """ return _ffi_api.SBijectiveLayoutForwardShape(self, shape) # type: ignore def backward_shape(self, shape): """Given the shape of the dst-layout, infer the src shape. Parameters ---------- shape: Array of Expr The shape in dst-layout. Returns ------- src_shape: Array of Expr The inferred shape in src-layout. """ return _ffi_api.SBijectiveLayoutBackwardShape(self, shape) # type: ignore def slayout(layout_str: str, dtype: str = "int32") -> SLayout: """Create a layout node from a string. Parameters ---------- layout_str : str A layout representation is composed of upper cases, lower cases and numbers, where upper case indicates a primal axis and the corresponding lower case with factor size indicates the subordinate axis. For example, NCHW16c can describe a 5-D tensor of [batch_size, channel, height, width, channel_block]. Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel). dtype : str The dtype of generated axes vars in the returned layout. It is required to be integer type. Returns ------- layout : SLayout The created layout """ return _ffi_api.SLayout(layout_str, dtype) # type: ignore def sbijective_layout(src_layout: str | SLayout, dst_layout: str | SLayout) -> SBijectiveLayout: """Create a bijective layout mapping. Parameters ---------- src_layout : str or SLayout source layout. dst_layout : str or SLayout destination layout. Returns ------- sbijective_layout : SBijectiveLayout The created bijective layout """ if isinstance(src_layout, str): src_layout = slayout(src_layout) if isinstance(dst_layout, str): dst_layout = slayout(dst_layout) return _ffi_api.SBijectiveLayout(src_layout, dst_layout) # type: ignore