google-research

1
// Copyright 2024 The Google Research Authors.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14

15
#include "sparse/ops/cc/transpose_launcher.h"
16

17
#include <limits>
18

19
namespace sgk {
20

21
void AllocateTransposeWorkspace(
22
    tensorflow::OpKernelContext *context, const Eigen::ThreadPoolDevice &d,
23
    int m, int n, int nonzeros, const float *values, const int *row_offsets,
24
    const int *column_indices, float *output_values, int *output_row_offsets,
25
    int *output_column_indices, tensorflow::Tensor *workspace) {
26
  // To transpose the matrix, we blow up the tensor into it's
27
  // dense, transposed representation and compress it back down.
28
  tensorflow::TensorShape shape = {m * n};
29
  OP_REQUIRES_OK(
30
      context, context->allocate_temp(tensorflow::DT_FLOAT, shape, workspace));
31
}
32

33
void LaunchTranspose(const Eigen::ThreadPoolDevice &d, int m, int n,
34
                     int nonzeros, const float *values, const int *row_offsets,
35
                     const int *column_indices, float *output_values,
36
                     int *output_row_offsets, int *output_column_indices,
37
                     float *workspace) {
38
  // Expand the tensor into it's tranposed dense representation.
39
  //
40
  // NOTE: We set the invalid values in the tensor to infinity. This
41
  // This avoids issues with the case where we have zero valued weights
42
  // in the sparse matrix.
43
  for (int i = 0; i < m * n; ++i) {
44
    workspace[i] = std::numeric_limits<float>::infinity();
45
  }
46
  for (int i = 0; i < m; ++i) {
47
    for (int l = row_offsets[i]; l < row_offsets[i + 1]; ++l) {
48
      int j = column_indices[l];
49
      workspace[j * m + i] = values[l];
50
    }
51
  }
52

53
  // Compress the matrix back down to it's sparse representation. Note
54
  // that the matrix is transposed, so 'n' is the number of rows and
55
  // 'm' is the number of columns.
56
  int offset = 0;
57
  output_row_offsets[0] = 0;
58
  for (int i = 0; i < n; ++i) {    // loop over rows.
59
    for (int j = 0; j < m; ++j) {  // loop over columns.
60
      int idx = i * m + j;
61
      if (workspace[idx] == std::numeric_limits<float>::infinity()) {
62
        continue;
63
      }
64
      DCHECK_LT(offset, nonzeros);
65
      output_values[offset] = workspace[idx];
66
      output_column_indices[offset] = j;
67
      ++offset;
68
    }
69
    output_row_offsets[i + 1] = offset;
70
  }
71
}
72

73
}  // namespace sgk
74