2 * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
3 * Copyright 2017 The TensorFlow Authors. All Rights Reserved.
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
9 * http://www.apache.org/licenses/LICENSE-2.0
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
18 #ifndef __NNFW_CKER_MATRIX_BAND_PART_H__
19 #define __NNFW_CKER_MATRIX_BAND_PART_H__
21 #include "cker/Shape.h"
30 void MatrixBandPart(const T num_lower_diags, const T num_upper_diags, const Shape &input_shape,
31 const float *input_data, const Shape &output_shape, float *output_data)
33 auto last_dim = input_shape.DimensionsCount() - 1;
36 for (int dim = 0; dim < last_dim - 2; dim++)
38 batch_num += input_shape.Dims(dim);
41 const T row_num = input_shape.Dims(last_dim - 1);
42 const T col_num = input_shape.Dims(last_dim);
44 if (!(num_lower_diags <= row_num))
45 throw std::runtime_error(
46 "MatrixBandPart : num_lower must be negative or less or equal to number of rows");
48 if (!(num_upper_diags <= col_num))
49 throw std::runtime_error(
50 "MatrixBandPart : num_upper must be negative or less or equal to number of columns");
52 std::fill(output_data, output_data + output_shape.FlatSize(), 0); // output matrix init
54 // reference code, without multithreading
55 for (T batch = 0; batch < batch_num; ++batch)
57 for (T row = 0; row < row_num; ++row)
59 auto output = output_data + (batch * row_num * col_num + row * col_num);
60 auto input = input_data + (batch * row_num * col_num + row * col_num);
63 num_lower_diags < 0 ? 0 : std::min(col_num, std::max(T{0}, row - num_lower_diags));
64 const T band_end = num_upper_diags < 0 ? col_num : std::min(static_cast<T>(col_num),
65 row + num_upper_diags + 1);
67 for (T band_idx = band_start; band_idx < band_end; band_idx++)
69 output[band_idx] = input[band_idx];
77 #endif // __NNFW_CKER_MATRIX_BAND_PART_H__