# 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, too-many-arguments, too-many-nested-blocks """Sparse_Reshape operator""" from tvm.script.ir_builder import IRBuilder from tvm.script.ir_builder import tirx as T from tvm.te import div, extern, floordiv, floormod from tvm.tirx import Cast, decl_buffer def sparse_reshape( sparse_indices, prev_shape, new_shape, new_sparse_indices_shape, new_shape_shape, ): """ Reshape a Sparse Tensor Parameters ---------- sparse_indices : te.Expr A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the number of sparse values and n_dim is the number of dimensions of the dense_shape prev_shape : te.Expr A 1-D tensor containing the previous shape of the dense tensor new_shape : te.Expr A 1-D tensor containing the new shape of the dense tensor Returns ------- result: te.Expr Output tensor. Examples -------- .. code-block:: python sparse_indices = [[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 2, 3]] prev_shape = [2, 3, 4] new_shape = [9, -1] new_sparse_indices, new_shape = topi.sparse_reshape( sparse_indices, prev_shape, new_shape) new_sparse_indices = [[0, 0], [0, 1], [1, 2], [4, 2], [8, 1]] new_shape = [9, 4] """ def gen_ir( sparse_indices_ptr, prev_shape_ptr, new_shape_ptr, new_sparse_indices_ptr, out_new_shape_ptr, ): with IRBuilder() as ib: sparse_indices = T.buffer_proxy(sparse_indices_ptr) prev_shape = T.buffer_proxy(prev_shape_ptr) new_shape = T.buffer_proxy(new_shape_ptr) out_new_shape = T.buffer_proxy(out_new_shape_ptr) new_sparse_indices = T.buffer_proxy(new_sparse_indices_ptr) prev_shape_size = prev_shape_ptr.shape[0] new_shape_size = new_shape_ptr.shape[0] multipliers_buf = T.alloc_buffer([prev_shape_size], new_shape_ptr.dtype, scope="local") multipliers = T.buffer_proxy(multipliers_buf) dividers_buf = T.alloc_buffer([new_shape_size], new_shape_ptr.dtype, scope="local") dividers = T.buffer_proxy(dividers_buf) flattened_indices_buf = T.alloc_buffer( [sparse_indices_ptr.shape[0]], new_shape_ptr.dtype, scope="local" ) flattened_indices = T.buffer_proxy(flattened_indices_buf) total_ele_buf = T.alloc_buffer([1], new_shape_ptr.dtype, scope="local") total_ele = T.buffer_proxy(total_ele_buf) division_total_ele_buf = T.alloc_buffer([1], new_shape_ptr.dtype, scope="local") division_total_ele = T.buffer_proxy(division_total_ele_buf) equal_shape_buf = T.alloc_buffer([1], "bool", scope="local") equal_shape = T.buffer_proxy(equal_shape_buf) total_ele[0] = prev_shape[0] # Cumulative Reverse Exclusive Multiply multipliers[prev_shape_size - 1] = Cast(new_shape_ptr.dtype, 1) with T.serial(0, prev_shape_size - 1) as i_: i = i_ + 1 multipliers[prev_shape_size - 1 - i] = ( prev_shape[prev_shape_size - i] * multipliers[prev_shape_size - i] ) total_ele[0] *= prev_shape[prev_shape_size - i] division_total_ele[0] = Cast(new_shape_ptr.dtype, 1) with T.serial(0, new_shape_size) as i: with T.If(new_shape[i] != -1): with T.Then(): division_total_ele[0] *= new_shape[i] # Compute true output shape (replace negative ones) with T.serial(0, new_shape_size) as i: with T.If(new_shape[i] == -1): with T.Then(): out_new_shape[i] = Cast( new_shape_ptr.dtype, div(total_ele[0], division_total_ele[0]) ) with T.Else(): out_new_shape[i] = new_shape[i] # Check if prev_shape and new_shape are equal equal_shape[0] = True with T.If(prev_shape_size == new_shape_size): with T.Then(): with T.serial(0, prev_shape_size) as i: with T.If(prev_shape[i] != out_new_shape[i]): with T.Then(): equal_shape[0] = False with T.Else(): equal_shape[0] = False # Return same inputs if shapes are equal with T.If(equal_shape[0]): with T.Then(): with T.parallel(0, sparse_indices_ptr.shape[0]) as i: with T.serial(0, sparse_indices_ptr.shape[1]) as j: new_sparse_indices[i, j] = sparse_indices[i, j] # Else compute new_sparse_indices with T.Else(): dividers[new_shape_size - 1] = Cast(new_shape_ptr.dtype, 1) with T.serial(0, new_shape_size - 1) as i_: i = i_ + 1 dividers[new_shape_size - 1 - i] = ( dividers[new_shape_size - i] * out_new_shape[new_shape_size - i] ) with T.parallel(0, sparse_indices_ptr.shape[0]) as i: flattened_indices[i] = Cast(new_shape_ptr.dtype, 0) with T.serial(0, sparse_indices_ptr.shape[1]) as j: flattened_indices[i] += sparse_indices[i, j] * multipliers[j] with T.parallel(0, new_sparse_indices_ptr.shape[0]) as i: current_element_buf = T.alloc_buffer( [1], new_shape_ptr.dtype, scope="local" ) current_element = T.buffer_proxy(current_element_buf) current_element[0] = flattened_indices[i] with T.serial(0, new_sparse_indices_ptr.shape[1]) as j: new_sparse_indices[i, j] = Cast( sparse_indices_ptr.dtype, floordiv(current_element[0], dividers[j]), ) current_element[0] = floormod(current_element[0], dividers[j]) return ib.get() new_sparse_indices_buf = decl_buffer( new_sparse_indices_shape, sparse_indices.dtype, "new_sparse_indices_buf" ) new_shape_buf = decl_buffer(new_shape_shape, prev_shape.dtype, "new_shape_buf") return extern( [new_sparse_indices_shape, new_shape_shape], [sparse_indices, prev_shape, new_shape], lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0], outs[1]), out_buffers=[new_sparse_indices_buf, new_shape_buf], name="sparse_reshape_cpu", tag="sparse_reshape_cpu", )