tensor-compiler.github.io/examples/spgemm_assembly.c at master · tensor-compiler/tensor-compiler.github.io · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
// Generated by the Tensor Algebra Compiler (tensor-compiler.org)
// taco "A(i,j)=B(i,k)*C(k,j)" -f=A:ds:0,1 -f=B:ds:0,1 -f=C:ds:0,1 -s="reorder(i,k,j)" -s="precompute(B(i,k)*C(k,j),j,j)" -s="assemble(A,Insert)" -s="parallelize(i,CPUThread,NoRaces)" -write-source=taco_kernel.c -write-compute=taco_compute.c -write-assembly=taco_assembly.c

int assemble(taco_tensor_t *A, taco_tensor_t *B, taco_tensor_t *C) {
  int A1_dimension = (int)(A->dimensions[0]);
  int* restrict A2_pos = (int*)(A->indices[1][0]);
  int* restrict A2_crd = (int*)(A->indices[1][1]);
  double* restrict A_vals = (double*)(A->vals);
  int B1_dimension = (int)(B->dimensions[0]);
  int* restrict B2_pos = (int*)(B->indices[1][0]);
  int* restrict B2_crd = (int*)(B->indices[1][1]);
  int C1_dimension = (int)(C->dimensions[0]);
  int C2_dimension = (int)(C->dimensions[1]);
  int* restrict C2_pos = (int*)(C->indices[1][0]);
  int* restrict C2_crd = (int*)(C->indices[1][1]);

  int32_t* restrict A2_nnz = 0;
  A2_nnz = (int32_t*)malloc(sizeof(int32_t) * B1_dimension);

  int32_t* restrict qworkspace_index_list_all = 0;
  qworkspace_index_list_all = (int32_t*)malloc(sizeof(int32_t) * (C2_dimension * omp_get_max_threads()));
  bool* restrict qworkspace_already_set_all = calloc((C2_dimension * omp_get_max_threads()), sizeof(bool));

  #pragma omp parallel for schedule(runtime)
  for (int32_t i = 0; i < B1_dimension; i++) {
    int32_t qworkspace_index_list_size = 0;
    int32_t* restrict qworkspace_index_list = qworkspace_index_list_all + C2_dimension * omp_get_thread_num();
    bool* restrict qworkspace_already_set = qworkspace_already_set_all + C2_dimension * omp_get_thread_num();
    for (int32_t kB = B2_pos[i]; kB < B2_pos[(i + 1)]; kB++) {
      int32_t k = B2_crd[kB];
      for (int32_t jC = C2_pos[k]; jC < C2_pos[(k + 1)]; jC++) {
        int32_t j = C2_crd[jC];
        if (!qworkspace_already_set[j]) {
          qworkspace_index_list[qworkspace_index_list_size] = j;
          qworkspace_already_set[j] = 1;
          qworkspace_index_list_size++;
        }
      }
    }
    int32_t tjA2_nnz_val = 0;
    for (int32_t qworkspace_index_locator = 0; qworkspace_index_locator < qworkspace_index_list_size; qworkspace_index_locator++) {
      int32_t j = qworkspace_index_list[qworkspace_index_locator];
      tjA2_nnz_val++;
      qworkspace_already_set[j] = 0;
    }
    A2_nnz[i] = tjA2_nnz_val;
  }

  free(qworkspace_index_list_all);
  free(qworkspace_already_set_all);

  A2_pos = (int32_t*)malloc(sizeof(int32_t) * (A1_dimension + 1));
  A2_pos[0] = 0;
  for (int32_t i = 0; i < A1_dimension; i++) {
    A2_pos[i + 1] = A2_pos[i] + A2_nnz[i];
  }
  A2_crd = (int32_t*)malloc(sizeof(int32_t) * A2_pos[A1_dimension]);
  A_vals = (double*)malloc(sizeof(double) * A2_pos[A1_dimension]);

  int32_t* restrict workspace_index_list_all = 0;
  workspace_index_list_all = (int32_t*)malloc(sizeof(int32_t) * (C2_dimension * omp_get_max_threads()));
  bool* restrict workspace_already_set_all = calloc((C2_dimension * omp_get_max_threads()), sizeof(bool));

  #pragma omp parallel for schedule(runtime)
  for (int32_t i = 0; i < B1_dimension; i++) {
    int32_t workspace_index_list_size = 0;
    int32_t* restrict workspace_index_list = workspace_index_list_all + C2_dimension * omp_get_thread_num();
    bool* restrict workspace_already_set = workspace_already_set_all + C2_dimension * omp_get_thread_num();
    for (int32_t kB = B2_pos[i]; kB < B2_pos[(i + 1)]; kB++) {
      int32_t k = B2_crd[kB];
      for (int32_t jC = C2_pos[k]; jC < C2_pos[(k + 1)]; jC++) {
        int32_t j = C2_crd[jC];
        if (!workspace_already_set[j]) {
          workspace_index_list[workspace_index_list_size] = j;
          workspace_already_set[j] = 1;
          workspace_index_list_size++;
        }
      }
    }
    qsort(workspace_index_list, workspace_index_list_size, sizeof(int32_t), cmp);

    for (int32_t workspace_index_locator = 0; workspace_index_locator < workspace_index_list_size; workspace_index_locator++) {
      int32_t j = workspace_index_list[workspace_index_locator];
      int32_t pA2 = A2_pos[i];
      A2_pos[i] = A2_pos[i] + 1;
      A2_crd[pA2] = j;
      workspace_already_set[j] = 0;
    }
  }

  free(workspace_index_list_all);
  free(workspace_already_set_all);

  for (int32_t p = 0; p < A1_dimension; p++) {
    A2_pos[A1_dimension - p] = A2_pos[((A1_dimension - p) - 1)];
  }
  A2_pos[0] = 0;

  free(A2_nnz);

  A->indices[1][0] = (uint8_t*)(A2_pos);
  A->indices[1][1] = (uint8_t*)(A2_crd);
  A->vals = (uint8_t*)A_vals;
  return 0;
}