|
- #!/usr/bin/env python3
- # coding: utf-8
- # Copyright 2019 Huawei Technologies Co., Ltd
- #
- # Licensed 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.
-
- """operator dsl function:argmin_argmax_common"""
- import akg.tvm
- import akg.topi
- from akg.lang import cce as dav
- from akg.utils import custom_tiling as ct_util, validation_check as vc_util
- from akg.utils.dsl_create import get_reduce_out_shape
- from akg.utils.format_transform import refine_reduce_axis, get_shape
- from akg.utils.dynamic_shape import shape_is_dynamic, set_dynamic_shape_limit_for_tensor
-
-
- def argminmax_tiling_strategy(out_shape, axis):
- """Custom tiling strategy for argminmax op."""
- strategy = list()
- # when reduce axis is one, we do not need any strategy
- if out_shape[axis] == 1:
- return strategy
-
- # if reduce first axis, it will transpose to last axis
- # so here we adapt to this change
- if axis == 0:
- temp = out_shape[0]
- out_shape = out_shape[1:]
- out_shape.append(temp)
- axis = len(out_shape) - 1
-
- # eliminate single axis, which will automatically disappear in halide ir
- # and adjust axis if it is influenced
- shrink = list()
- for i, shp in enumerate(out_shape):
- if shp == 1:
- if i < axis:
- axis -= 1
- else:
- shrink.append(shp)
-
- for i, _ in enumerate(shrink):
- if i == axis:
- strategy.append(ct_util.create_constraint_on_axis(
- values="FULL",
- constraints=ct_util.TileConstraint.MAX,
- axis=i)[0])
- else:
- strategy.append(ct_util.create_constraint_on_axis(
- values=1,
- constraints=ct_util.TileConstraint.FACTOR,
- axis=i)[0])
- return strategy
-
-
- @vc_util.check_input_type(akg.tvm.tensor.Tensor, int, (str, type(None)))
- def common(data, axis, method="min"):
- """
- Returns the index with the max or min value across axes of a tensor.
-
- Note:
- method can be "max" or "min" to get argmax or argmin.
-
- Args:
- data (tvm.tensor.Tensor): Tensor of type float16, float32, int8, int32.
- axis (int): Describe the axis of input tensor.
- method (str): Can be "max" or "min".
-
- Returns:
- tvm.tensor.Tensor, has type of int32.
- """
- shape = get_shape(data)
- dtype = data.dtype
-
- vc_util.ops_dtype_check(data.dtype, [vc_util.DtypeForDavinci.ALL_FLOAT, vc_util.DtypeForDavinci.ALL_INT])
- vc_util.reduce_axis_check(shape, axis)
- real_axis = refine_reduce_axis(shape, axis)[0]
- out_shape = get_reduce_out_shape(shape, axis=axis)
- attr_map = {}
- if shape_is_dynamic(data):
- attr_map["dynamic_shape"] = set_dynamic_shape_limit_for_tensor(data, 4096, real_axis)
- if dtype != "float16":
- data = akg.topi.cast(data, "float16")
- k = akg.tvm.reduce_axis((0, data.shape[real_axis]), "k")
- if axis in (len(shape) - 1, -1):
- if method == "min":
- reducer = akg.tvm.comm_reducer(
- lambda x, y: dav.fargmin(x, y), lambda t: akg.tvm.max_value(t))
- elif method == "max":
- reducer = akg.tvm.comm_reducer(
- lambda x, y: dav.fargmax(x, y), lambda t: akg.tvm.min_value(t))
- else:
- raise ValueError("not support " + method)
-
- if len(data.shape) == 1:
- res = akg.tvm.compute((1,), lambda i: reducer(data[k], axis=k))
- else:
- res = akg.tvm.compute(out_shape,
- lambda *indice:
- reducer(data(*indice, k), axis=k))
-
- res = akg.tvm.compute(out_shape,
- lambda *indice: res(*indice).astype("int32"),
- "argred_output")
- elif axis in (0, -len(shape)):
- tmp_idx = akg.tvm.compute(shape[1:],
- lambda *indice: akg.tvm.const(0.0, "float16"),
- name='tmp_index')
- local_data = akg.tvm.compute(shape[1:],
- lambda *indice: data(0, *indice),
- name="tmp_data")
- for idx in range(shape[axis] - 1):
- if method == 'min':
- tmp_idx = akg.tvm.compute(
- shape[1:],
- lambda *indice, ite_idx=idx:
- akg.tvm.expr.Select(
- local_data(*indice) > data(ite_idx + 1, *indice),
- akg.tvm.const(ite_idx + 1, "float16"),
- tmp_idx(*indice)
- ))
- local_data = akg.tvm.compute(
- shape[1:],
- lambda *indice, ite_idx=idx:
- akg.tvm.expr.Select(
- local_data(*indice) > data(ite_idx + 1, *indice),
- data(ite_idx + 1, *indice),
- local_data(*indice)
- ))
- elif method == "max":
- tmp_idx = akg.tvm.compute(
- shape[1:],
- lambda *indice, ite_idx=idx:
- akg.tvm.expr.Select(
- local_data(*indice) < data(ite_idx + 1, *indice),
- akg.tvm.const(ite_idx + 1, "float16"),
- tmp_idx(*indice)
- ))
- local_data = akg.tvm.compute(
- shape[1:],
- lambda *indice, ite_idx=idx:
- akg.tvm.expr.Select(
- local_data(*indice) < data(ite_idx + 1, *indice),
- data(ite_idx + 1, *indice),
- local_data(*indice)
- ))
- else:
- raise ValueError("not support " + method)
-
- res = akg.tvm.compute(out_shape,
- lambda *indice: tmp_idx(*indice).astype("int32"),
- "cast1")
- else:
- raise ValueError("Argmax only support first axis and is last axis now!")
-
- lager = out_shape if len(out_shape) > len(shape) else shape
- strategy = argminmax_tiling_strategy(lager, real_axis)
- if strategy:
- attr_map["custom_tiling"] = strategy
- return res, attr_map
|
