summation.cxx 71.7 KB
Newer Older
1
#include "summation.h"
2 3 4 5 6 7 8 9
#include "../scaling/strp_tsr.h"
#include "../mapping/mapping.h"
#include "../mapping/distribution.h"
#include "../tensor/untyped_tensor.h"
#include "../shared/util.h"
#include "../shared/memcontrol.h"
#include "sym_seq_sum.h"
#include "sum_tsr.h"
10 11
#include "../symmetry/sym_indices.h"
#include "../symmetry/symmetrization.h"
12
#include "../redistribution/nosym_transp.h"
13
#include "../redistribution/redist.h"
solomon's avatar
solomon committed
14
#include "../scaling/scaling.h"
15 16 17

namespace CTF_int {

18 19
  using namespace CTF;

20
  summation::~summation(){
solomon's avatar
solomon committed
21 22
    if (idx_A != NULL) cdealloc(idx_A);
    if (idx_B != NULL) cdealloc(idx_B);
23 24
  }

solomon's avatar
solomon committed
25
  summation::summation(summation const & other){
26
    A     = other.A;
27
    idx_A = (int*)alloc(sizeof(int)*other.A->order);
28
    memcpy(idx_A, other.idx_A, sizeof(int)*other.A->order);
29
    B     = other.B;
30
    idx_B = (int*)alloc(sizeof(int)*other.B->order);
31
    memcpy(idx_B, other.idx_B, sizeof(int)*other.B->order);
solomon's avatar
solomon committed
32
    if (other.is_custom){
33
      func      = other.func;
solomon's avatar
solomon committed
34
      is_custom = 1;
35 36 37
    } else is_custom = 0; 
    alpha = other.alpha;
    beta  = other.beta;
solomon's avatar
solomon committed
38 39
  }

40 41 42 43 44 45 46 47 48
  summation::summation(tensor *     A_,
                       int const *  idx_A_,
                       char const * alpha_,
                       tensor *     B_,
                       int const *  idx_B_,
                       char const * beta_){
    A         = A_;
    alpha     = alpha_;
    B         = B_;
49 50 51
    beta      = beta_;
    is_custom = 0;

52 53
    idx_A     = (int*)alloc(sizeof(int)*A->order);
    idx_B     = (int*)alloc(sizeof(int)*B->order);
54 55

    memcpy(idx_A, idx_A_, sizeof(int)*A->order);
56
    memcpy(idx_B, idx_B_, sizeof(int)*B->order);
57 58 59 60 61 62 63 64 65 66 67
  }

  summation::summation(tensor *     A_,
                       char const * cidx_A,
                       char const * alpha_,
                       tensor *     B_,
                       char const * cidx_B,
                       char const * beta_){
    A         = A_;
    alpha     = alpha_;
    B         = B_;
68
    beta      = beta_;
69
    is_custom = 0;
70 71
    
    conv_idx(A->order, cidx_A, &idx_A, B->order, cidx_B, &idx_B);
72
  }
73

74
 
75 76 77 78 79 80 81
  summation::summation(tensor *                A_,
                       int const *             idx_A_,
                       char const *            alpha_,
                       tensor *                B_,
                       int const *             idx_B_,
                       char const *            beta_,
                       univar_function const * func_){
82
    A         = A_;
83
    alpha     = alpha_;
84
    B         = B_;
85 86 87 88
    beta      = beta_;
    func      = func_;
    is_custom = 1;

89 90
    idx_A     = (int*)alloc(sizeof(int)*A->order);
    idx_B     = (int*)alloc(sizeof(int)*B->order);
91 92

    memcpy(idx_A, idx_A_, sizeof(int)*A->order);
93
    memcpy(idx_B, idx_B_, sizeof(int)*B->order);
94 95 96
  }

 
97 98 99 100 101 102 103
  summation::summation(tensor *                A_,
                       char const *            cidx_A,
                       char const *            alpha_,
                       tensor *                B_,
                       char const *            cidx_B,
                       char const *            beta_,
                       univar_function const * func_){
104 105 106
    A         = A_;
    alpha     = alpha_;
    B         = B_;
107
    beta      = beta_;
108
    func      = func_;
109
    is_custom = 1;
110 111

    conv_idx(A->order, cidx_A, &idx_A, B->order, cidx_B, &idx_B);
112 113
  }

114
  void summation::execute(bool run_diag){
115 116 117 118
#if DEBUG >= 2
    if (A->wrld->cdt.rank == 0) printf("Summation::execute (head):\n");
    print();
#endif
119
    if (A->is_sparse || B->is_sparse){
120 121
      int stat = sym_sum_tsr(run_diag);
      assert(stat == SUCCESS); 
122 123 124 125
    } else {
      int stat = home_sum_tsr(run_diag);
      assert(stat == SUCCESS); 
    }
126 127 128 129 130 131 132
  }
  
  double summation::estimate_time(){
    assert(0); //FIXME
    return 0.0;
  }

133 134
  void summation::get_fold_indices(int *  num_fold,
                                   int ** fold_idx){
135 136 137
    int i, in, num_tot, nfold, broken;
    int iA, iB, inA, inB, iiA, iiB;
    int * idx_arr, * idx;
138 139
    inv_idx(A->order, idx_A,
            B->order, idx_B,
140
            &num_tot, &idx_arr);
141
    CTF_int::alloc_ptr(num_tot*sizeof(int), (void**)&idx);
142 143 144 145 146

    for (i=0; i<num_tot; i++){
      idx[i] = 1;
    }
    
147
    for (iA=0; iA<A->order; iA++){
148 149
      i      = idx_A[iA];
      iB     = idx_arr[2*i+1];
150
      broken = 0;
151
      inA    = iA;
152
      do {
153
        in = idx_A[inA];
154 155 156
        inB = idx_arr[2*in+1];
        if (((inA>=0) + (inB>=0) != 2) ||
            (iB != -1 && inB - iB != in-i) ||
157
            (iB != -1 && A->sym[inA] != B->sym[inB])){
158
          broken = 1;
159
          //printf("index in = %d inA = %d inB = %d is broken symA = %d symB = %d\n",in, inA, inB, A->sym[inA], B->sym[inB]);
160 161
        }
        inA++;
162
      } while (A->sym[inA-1] != NS);
163 164
      if (broken){
        for (iiA=iA;iiA<inA;iiA++){
165
          idx[idx_A[iiA]] = 0;
166 167 168 169
        }
      }
    }

170
    for (iB=0; iB<B->order; iB++){
171 172
      i      = idx_B[iB];
      iA     = idx_arr[2*i+0];
173
      broken = 0;
174
      inB    = iB;
175
      do {
176
        in = idx_B[inB];
177 178 179
        inA = idx_arr[2*in+0];
        if (((inB>=0) + (inA>=0) != 2) ||
            (iA != -1 && inA - iA != in-i) ||
180
            (iA != -1 && B->sym[inB] != A->sym[inA])){
181 182 183
          broken = 1;
        }
        inB++;
184
      } while (B->sym[inB-1] != NS);
185 186
      if (broken){
        for (iiB=iB;iiB<inB;iiB++){
187
          idx[idx_B[iiB]] = 0;
188 189 190 191 192 193 194 195 196 197 198 199 200 201
        }
      }
    }
    

    nfold = 0;
    for (i=0; i<num_tot; i++){
      if (idx[i] == 1){
        idx[nfold] = i;
        nfold++;
      }
    }
    *num_fold = nfold;
    *fold_idx = idx;
solomon's avatar
solomon committed
202
    CTF_int::cdealloc(idx_arr);
203 204 205 206
  }

  int summation::can_fold(){
    int i, j, nfold, * fold_idx;
207
    //FIXME: fold sparse tensors into CSR form
208
    if (A->is_sparse || B->is_sparse) return 0;
209

210 211 212
    for (i=0; i<A->order; i++){
      for (j=i+1; j<A->order; j++){
        if (idx_A[i] == idx_A[j]) return 0;
213 214
      }
    }
215 216 217
    for (i=0; i<B->order; i++){
      for (j=i+1; j<B->order; j++){
        if (idx_B[i] == idx_B[j]) return 0;
218 219
      }
    }
220
    get_fold_indices(&nfold, &fold_idx);
solomon's avatar
solomon committed
221
    CTF_int::cdealloc(fold_idx);
222 223 224
    /* FIXME: 1 folded index is good enough for now, in the future model */
    return nfold > 0;
  }
225 226 227 228 229 230 231
 
  void summation::get_fold_sum(summation *& fold_sum,
                               int &        all_fdim_A,
                               int &        all_fdim_B,
                               int *&       all_flen_A,
                               int *&       all_flen_B){
    int i, j, nfold, nf;
232 233 234 235 236
    int * fold_idx, * fidx_map_A, * fidx_map_B;
    tensor * ftsr_A, * ftsr_B;

    get_fold_indices(&nfold, &fold_idx);
    if (nfold == 0){
solomon's avatar
solomon committed
237
      CTF_int::cdealloc(fold_idx);
238
      assert(0);
239 240 241
    }

    /* overestimate this space to not bother with it later */
242 243
    CTF_int::alloc_ptr(nfold*sizeof(int), (void**)&fidx_map_A);
    CTF_int::alloc_ptr(nfold*sizeof(int), (void**)&fidx_map_B);
244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272


    A->fold(nfold, fold_idx, idx_A, 
            &all_fdim_A, &all_flen_A);
    B->fold(nfold, fold_idx, idx_B, 
            &all_fdim_B, &all_flen_B);

    nf = 0;
    for (i=0; i<A->order; i++){
      for (j=0; j<nfold; j++){
        if (A->sym[i] == NS && idx_A[i] == fold_idx[j]){
          fidx_map_A[nf] = j;
          nf++;
        }
      }
    }
    nf = 0;
    for (i=0; i<B->order; i++){
      for (j=0; j<nfold; j++){
        if (B->sym[i] == NS && idx_B[i] == fold_idx[j]){
          fidx_map_B[nf] = j;
          nf++;
        }
      }
    }

    ftsr_A = A->rec_tsr;
    ftsr_B = B->rec_tsr;

273
    int * sidx_A, * sidx_B;
274
    CTF_int::conv_idx<int>(ftsr_A->order, fidx_map_A, &sidx_A,
Edgar Solomonik's avatar
Edgar Solomonik committed
275 276 277 278 279
                           ftsr_B->order, fidx_map_B, &sidx_B);
    
    cdealloc(fidx_map_A);
    cdealloc(fidx_map_B);
    cdealloc(fold_idx);
280

281
    fold_sum = new summation(A->rec_tsr, sidx_A, alpha, B->rec_tsr, sidx_B, beta);
solomon's avatar
solomon committed
282 283
    cdealloc(sidx_A);
    cdealloc(sidx_B);
284 285 286 287 288 289 290 291 292 293 294
  }

  int summation::map_fold(){
    int i, all_fdim_A, all_fdim_B;
    int nvirt_A, nvirt_B;
    int * fnew_ord_A, * fnew_ord_B;
    int * all_flen_A, * all_flen_B;
    int inr_stride;

    summation * fold_sum;
    get_fold_sum(fold_sum, all_fdim_A, all_fdim_B, all_flen_A, all_flen_B);
295
  #if DEBUG>=2
296
    if (A->wrld->rank == 0){
297 298
      printf("Folded summation type:\n");
    }
299
    fold_sum->print();//print_sum(&fold_type,0.0,0.0);
300 301 302
  #endif
   
    //for type order 1 to 3 
303 304 305
    fold_sum->get_len_ordering(&fnew_ord_A, &fnew_ord_B); 
    permute_target(fold_sum->A->order, fnew_ord_A, A->inner_ordering);
    permute_target(fold_sum->B->order, fnew_ord_B, B->inner_ordering);
306 307
    

solomon's avatar
solomon committed
308
    nvirt_A = A->calc_nvirt();
309
    for (i=0; i<nvirt_A; i++){
solomon's avatar
solomon committed
310
      nosym_transpose(all_fdim_A, A->inner_ordering, all_flen_A, 
311
                      A->data + A->sr->el_size*i*(A->size/nvirt_A), 1, A->sr);
312
    }
solomon's avatar
solomon committed
313
    nvirt_B = B->calc_nvirt();
314
    for (i=0; i<nvirt_B; i++){
solomon's avatar
solomon committed
315
      nosym_transpose(all_fdim_B, B->inner_ordering, all_flen_B, 
316
                      B->data + B->sr->el_size*i*(B->size/nvirt_B), 1, B->sr);
317 318 319
    }

    inr_stride = 1;
320 321
    for (i=0; i<fold_sum->A->order; i++){
      inr_stride *= fold_sum->A->pad_edge_len[i];
322 323
    }

solomon's avatar
solomon committed
324 325 326 327
    CTF_int::cdealloc(fnew_ord_A);
    CTF_int::cdealloc(fnew_ord_B);
    CTF_int::cdealloc(all_flen_A);
    CTF_int::cdealloc(all_flen_B);
328

Edgar Solomonik's avatar
Edgar Solomonik committed
329 330
    delete fold_sum;

solomon's avatar
solomon committed
331
    return inr_stride; 
332 333
  }

334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361
  double summation::est_time_fold(){
    int all_fdim_A, all_fdim_B;
    int * fnew_ord_A, * fnew_ord_B;
    int * all_flen_A, * all_flen_B;
    int * tAiord, * tBiord;

    summation * fold_sum;
    get_fold_sum(fold_sum, all_fdim_A, all_fdim_B, all_flen_A, all_flen_B);
    fold_sum->get_len_ordering(&fnew_ord_A, &fnew_ord_B); 
    CTF_int::alloc_ptr(all_fdim_A*sizeof(int), (void**)&tAiord);
    CTF_int::alloc_ptr(all_fdim_B*sizeof(int), (void**)&tBiord);
    memcpy(tAiord, A->inner_ordering, all_fdim_A*sizeof(int));
    memcpy(tBiord, B->inner_ordering, all_fdim_B*sizeof(int));

    permute_target(fold_sum->A->order, fnew_ord_A, tAiord);
    permute_target(fold_sum->B->order, fnew_ord_B, tBiord);
  
    A->is_folded = 0; 
    delete A->rec_tsr; 
    cdealloc(A->inner_ordering); 
    B->is_folded = 0; 
    delete B->rec_tsr; 
    cdealloc(B->inner_ordering); 

    double esttime = 0.0;

    esttime += A->calc_nvirt()*est_time_transp(all_fdim_A, tAiord, all_flen_A, 1, A->sr);
    esttime += 2.*B->calc_nvirt()*est_time_transp(all_fdim_B, tBiord, all_flen_B, 1, B->sr);
Edgar Solomonik's avatar
Edgar Solomonik committed
362 363 364 365 366 367 368 369 370

    delete fold_sum;

    cdealloc(all_flen_A);
    cdealloc(all_flen_B);
    cdealloc(tAiord);
    cdealloc(tBiord);
    cdealloc(fnew_ord_A);
    cdealloc(fnew_ord_B);
371 372 373 374 375
    return esttime;
  }


 
376 377
  void summation::get_len_ordering(int ** new_ordering_A,
                                   int ** new_ordering_B){
378
    int num_tot;
379 380
    int * ordering_A, * ordering_B, * idx_arr;
    
381 382
    CTF_int::alloc_ptr(sizeof(int)*A->order, (void**)&ordering_A);
    CTF_int::alloc_ptr(sizeof(int)*B->order, (void**)&ordering_B);
383

384 385
    inv_idx(A->order, idx_A,
            B->order, idx_B,
386
            &num_tot, &idx_arr);
387 388 389
    ASSERT(num_tot == A->order);
    ASSERT(num_tot == B->order);
    /*for (i=0; i<num_tot; i++){
390 391
      ordering_A[i] = idx_arr[2*i];
      ordering_B[i] = idx_arr[2*i+1];
392 393 394 395 396 397 398
    }*/
    for (int i=0; i<num_tot; i++){
      ordering_B[i] = i;
      for (int j=0; j<num_tot; j++){
        if (idx_A[j] == idx_B[i])
          ordering_A[i] = j;
      }
399
    }
solomon's avatar
solomon committed
400
    CTF_int::cdealloc(idx_arr);
401 402 403 404
    *new_ordering_A = ordering_A;
    *new_ordering_B = ordering_B;
  }

405 406 407 408 409 410 411
  tsum * summation::construct_sum(int inner_stride){
    int nvirt, i, iA, iB, order_tot, is_top, sA, sB, need_rep, i_A, i_B, j, k;
    int64_t blk_sz_A, blk_sz_B, vrt_sz_A, vrt_sz_B;
    int nphys_dim;
    int * idx_arr, * virt_dim, * phys_mapped;
    int * virt_blk_len_A, * virt_blk_len_B;
    int * blk_len_A, * blk_len_B;
solomon's avatar
solomon committed
412
    tsum * htsum = NULL , ** rec_tsum = NULL;
413
    mapping * map;
solomon's avatar
solomon committed
414
    strp_tsr * str_A, * str_B;
415 416

    is_top = 1;
solomon's avatar
solomon committed
417 418
    inv_idx(A->order, idx_A,
            B->order, idx_B,
419 420
            &order_tot, &idx_arr);

solomon's avatar
solomon committed
421
    nphys_dim = A->topo->order;
422

423 424 425 426 427 428
    CTF_int::alloc_ptr(sizeof(int)*A->order,    (void**)&blk_len_A);
    CTF_int::alloc_ptr(sizeof(int)*B->order,    (void**)&blk_len_B);
    CTF_int::alloc_ptr(sizeof(int)*A->order,    (void**)&virt_blk_len_A);
    CTF_int::alloc_ptr(sizeof(int)*B->order,    (void**)&virt_blk_len_B);
    CTF_int::alloc_ptr(sizeof(int)*order_tot,   (void**)&virt_dim);
    CTF_int::alloc_ptr(sizeof(int)*nphys_dim*2, (void**)&phys_mapped);
429 430 431 432
    memset(phys_mapped, 0, sizeof(int)*nphys_dim*2);


    /* Determine the block dimensions of each local subtensor */
solomon's avatar
solomon committed
433 434 435
    blk_sz_A = A->size;
    blk_sz_B = B->size;
    calc_dim(A->order, blk_sz_A, A->pad_edge_len, A->edge_map,
436
             &vrt_sz_A, virt_blk_len_A, blk_len_A);
solomon's avatar
solomon committed
437
    calc_dim(B->order, blk_sz_B, B->pad_edge_len, B->edge_map,
438 439 440
             &vrt_sz_B, virt_blk_len_B, blk_len_B);

    /* Strip out the relevant part of the tensor if we are contracting over diagonal */
solomon's avatar
solomon committed
441 442
    sA = strip_diag(A->order, order_tot, idx_A, vrt_sz_A,
                           A->edge_map, A->topo, A->sr,
443
                           blk_len_A, &blk_sz_A, &str_A);
solomon's avatar
solomon committed
444 445
    sB = strip_diag(B->order, order_tot, idx_B, vrt_sz_B,
                           B->edge_map, B->topo, B->sr,
446 447
                           blk_len_B, &blk_sz_B, &str_B);
    if (sA || sB){
solomon's avatar
solomon committed
448
      if (A->wrld->cdt.rank == 0)
449
        DPRINTF(1,"Stripping tensor\n");
450
      strp_sum * ssum = new strp_sum(this);
solomon's avatar
solomon committed
451 452
      ssum->sr_A = A->sr;
      ssum->sr_B = B->sr;
453 454 455 456 457 458 459 460 461 462 463 464 465 466 467
      htsum = ssum;
      is_top = 0;
      rec_tsum = &ssum->rec_tsum;

      ssum->rec_strp_A = str_A;
      ssum->rec_strp_B = str_B;
      ssum->strip_A = sA;
      ssum->strip_B = sB;
    }

    nvirt = 1;
    for (i=0; i<order_tot; i++){
      iA = idx_arr[2*i];
      iB = idx_arr[2*i+1];
      if (iA != -1){
solomon's avatar
solomon committed
468
        map = &A->edge_map[iA];
469 470 471 472 473 474 475 476
        while (map->has_child) map = map->child;
        if (map->type == VIRTUAL_MAP){
          virt_dim[i] = map->np;
          if (sA) virt_dim[i] = virt_dim[i]/str_A->strip_dim[iA];
        }
        else virt_dim[i] = 1;
      } else {
        ASSERT(iB!=-1);
solomon's avatar
solomon committed
477
        map = &B->edge_map[iB];
478 479 480 481 482 483 484 485 486 487
        while (map->has_child) map = map->child;
        if (map->type == VIRTUAL_MAP){
          virt_dim[i] = map->np;
          if (sB) virt_dim[i] = virt_dim[i]/str_B->strip_dim[iA];
        }
        else virt_dim[i] = 1;
      }
      nvirt *= virt_dim[i];
    }

solomon's avatar
solomon committed
488 489
    for (i=0; i<A->order; i++){
      map = &A->edge_map[i];
490 491 492 493 494 495 496 497 498 499
      if (map->type == PHYSICAL_MAP){
        phys_mapped[2*map->cdt+0] = 1;
      }
      while (map->has_child) {
        map = map->child;
        if (map->type == PHYSICAL_MAP){
          phys_mapped[2*map->cdt+0] = 1;
        }
      }
    }
solomon's avatar
solomon committed
500 501
    for (i=0; i<B->order; i++){
      map = &B->edge_map[i];
502 503 504 505 506 507 508 509 510 511
      if (map->type == PHYSICAL_MAP){
        phys_mapped[2*map->cdt+1] = 1;
      }
      while (map->has_child) {
        map = map->child;
        if (map->type == PHYSICAL_MAP){
          phys_mapped[2*map->cdt+1] = 1;
        }
      }
    }
512 513 514 515 516
    bool need_perm = false;
    if (A->is_sparse || B->is_sparse){
      for (int i=0; i<A->order; i++){
        if (idx_arr[2*idx_A[i]+1] != i) need_perm = true;
      }
solomon's avatar
solomon committed
517 518 519
      for (int i=0; i<B->order; i++){
        if (idx_arr[2*idx_B[i]+1] != i) need_perm = true;
      }
520 521
    }
    if (need_perm){
solomon's avatar
solomon committed
522 523 524 525 526 527 528 529 530 531 532 533
      if (A->is_sparse){
        tsum_sp_pin_keys * sksum = new tsum_sp_pin_keys(this, 1);
        if (is_top){
          htsum = sksum;
          is_top = 0;
        } else {
          *rec_tsum = sksum;
        }
        rec_tsum = &sksum->rec_tsum;


        tsum_sp_permute * pmsum = new tsum_sp_permute(this, 1, virt_blk_len_A);
534
        *rec_tsum = pmsum;
solomon's avatar
solomon committed
535
        rec_tsum = &pmsum->rec_tsum;
536
      }
solomon's avatar
solomon committed
537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552
      if (B->is_sparse){
        tsum_sp_pin_keys * sksum = new tsum_sp_pin_keys(this, 0);
        if (is_top){
          htsum = sksum;
          is_top = 0;
        } else {
          *rec_tsum = sksum;
        }
        rec_tsum = &sksum->rec_tsum;


        tsum_sp_permute * pmsum = new tsum_sp_permute(this, 0, virt_blk_len_B);
        *rec_tsum = pmsum;
        rec_tsum = &pmsum->rec_tsum;
      }

553 554
    }

solomon's avatar
solomon committed
555
/*    bool need_sp_map = false;
556 557
    if (A->is_sparse || B->is_sparse){
      for (int i=0; i<B->order; i++){
558
        bool found_match = false;
559 560
        for (int j=0; j<A->order; j++){
          if (idx_B[i] == idx_A[j]) found_match = true;
561 562
        }
        if (!found_match) need_sp_map = true;
563 564 565
      }
    }

566 567 568 569 570 571 572 573 574
    if (need_sp_map){
      tsum_sp_map * smsum = new tsum_sp_map(this);
      if (is_top){
        htsum = smsum;
        is_top = 0;
      } else {
        *rec_tsum = smsum;
      }
      rec_tsum = &smsum->rec_tsum;
solomon's avatar
solomon committed
575
    }*/
576 577


578 579 580 581 582 583 584 585 586
    need_rep = 0;
    for (i=0; i<nphys_dim; i++){
      if (phys_mapped[2*i+0] == 0 ||
          phys_mapped[2*i+1] == 0){
        need_rep = 1;
        break;
      }
    }
    if (need_rep){
solomon's avatar
solomon committed
587
      if (A->wrld->cdt.rank == 0)
588 589
        DPRINTF(1,"Replicating tensor\n");

590 591
      tsum_replicate * rtsum = new tsum_replicate(this);
/*      rtsum->sr_A = A->sr;
solomon's avatar
solomon committed
592
      rtsum->sr_B = B->sr;
593 594 595
      rtsum->is_sparse_A = A->is_sparse;
      rtsum->is_sparse_B = B->is_sparse;
      rtsum->nnz_A = A->nnz_loc;
596
      rtsum->nnz_B = B->nnz_loc;*/
597 598 599 600 601 602
      if (is_top){
        htsum = rtsum;
        is_top = 0;
      } else {
        *rec_tsum = rtsum;
      }
603
      rec_tsum      = &rtsum->rec_tsum;
604 605 606 607
      rtsum->ncdt_A = 0;
      rtsum->ncdt_B = 0;
      rtsum->size_A = blk_sz_A;
      rtsum->size_B = blk_sz_B;
608 609
      rtsum->cdt_A  = NULL;
      rtsum->cdt_B  = NULL;
610 611 612 613 614 615 616 617 618
      for (i=0; i<nphys_dim; i++){
        if (phys_mapped[2*i+0] == 0 && phys_mapped[2*i+1] == 1){
          rtsum->ncdt_A++;
        }
        if (phys_mapped[2*i+1] == 0 && phys_mapped[2*i+0] == 1){
          rtsum->ncdt_B++;
        }
      }
      if (rtsum->ncdt_A > 0)
619
        CTF_int::alloc_ptr(sizeof(CommData*)*rtsum->ncdt_A, (void**)&rtsum->cdt_A);
620
      if (rtsum->ncdt_B > 0)
621
        CTF_int::alloc_ptr(sizeof(CommData*)*rtsum->ncdt_B, (void**)&rtsum->cdt_B);
622 623 624 625
      rtsum->ncdt_A = 0;
      rtsum->ncdt_B = 0;
      for (i=0; i<nphys_dim; i++){
        if (phys_mapped[2*i+0] == 0 && phys_mapped[2*i+1] == 1){
626 627 628
          rtsum->cdt_A[rtsum->ncdt_A] = &A->topo->dim_comm[i];
/*          if (rtsum->cdt_A[rtsum->ncdt_A].alive == 0)
            rtsum->cdt_A[rtsum->ncdt_A].activate(A->wrld->comm);*/
629 630 631
          rtsum->ncdt_A++;
        }
        if (phys_mapped[2*i+1] == 0 && phys_mapped[2*i+0] == 1){
632 633 634
          rtsum->cdt_B[rtsum->ncdt_B] = &B->topo->dim_comm[i];
/*          if (rtsum->cdt_B[rtsum->ncdt_B].alive == 0)
            rtsum->cdt_B[rtsum->ncdt_B].activate(B->wrld->comm);*/
635 636 637 638 639 640 641
          rtsum->ncdt_B++;
        }
      }
      ASSERT(rtsum->ncdt_A == 0 || rtsum->cdt_B == 0);
    }

    int * new_sym_A, * new_sym_B;
642
    CTF_int::alloc_ptr(sizeof(int)*A->order, (void**)&new_sym_A);
solomon's avatar
solomon committed
643
    memcpy(new_sym_A, A->sym, sizeof(int)*A->order);
644
    CTF_int::alloc_ptr(sizeof(int)*B->order, (void**)&new_sym_B);
solomon's avatar
solomon committed
645
    memcpy(new_sym_B, B->sym, sizeof(int)*B->order);
646 647

    /* Multiply over virtual sub-blocks */
648
    if (true){ //nvirt > 1){
649 650
      tsum_virt * tsumv = new tsum_virt(this);
/*      tsumv->sr_A = A->sr;
solomon's avatar
solomon committed
651
      tsumv->sr_B = B->sr;
652 653 654 655
      tsumv->is_sparse_A = A->is_sparse;
      tsumv->is_sparse_B = B->is_sparse;
      tsumv->nnz_A = A->nnz_loc;
      tsumv->nnz_B = B->nnz_loc;
656
      tsumv->nnz_blk_B = B->nnz_blk;*/
657 658 659 660 661 662
      if (is_top) {
        htsum = tsumv;
        is_top = 0;
      } else {
        *rec_tsum = tsumv;
      }
663 664 665 666
      rec_tsum         = &tsumv->rec_tsum;

      tsumv->num_dim   = order_tot;
      tsumv->virt_dim  = virt_dim;
667
//      tsumv->order_A   = A->order;
668
      tsumv->blk_sz_A  = vrt_sz_A;
669 670
//      tsumv->idx_map_A = idx_A;
//      tsumv->order_B   = B->order;
671
      tsumv->blk_sz_B  = vrt_sz_B;
672
//      tsumv->idx_map_B = idx_B;
673
      tsumv->buffer    = NULL;
solomon's avatar
solomon committed
674
    } else CTF_int::cdealloc(virt_dim);
675

676 677
    seq_tsr_sum * tsumseq = new seq_tsr_sum(this);
/*    tsumseq->sr_A = A->sr;
solomon's avatar
solomon committed
678
    tsumseq->sr_B = B->sr;
679 680 681
    tsumseq->is_sparse_A = A->is_sparse;
    tsumseq->is_sparse_B = B->is_sparse;
    tsumseq->nnz_A = A->nnz_loc;
682
    tsumseq->nnz_B = B->nnz_loc;*/
683 684 685 686 687
    if (inner_stride == -1){
      tsumseq->is_inner = 0;
    } else {
      tsumseq->is_inner = 1;
      tsumseq->inr_stride = inner_stride;
solomon's avatar
solomon committed
688 689
      tensor * itsr;
      itsr = A->rec_tsr;
690
      i_A = 0;
solomon's avatar
solomon committed
691 692
      for (i=0; i<A->order; i++){
        if (A->sym[i] == NS){
693
          for (j=0; j<itsr->order; j++){
solomon's avatar
solomon committed
694
            if (A->inner_ordering[j] == i_A){
695 696 697
              j=i;
              do {
                j--;
solomon's avatar
solomon committed
698
              } while (j>=0 && A->sym[j] != NS);
699 700 701 702 703 704 705 706 707 708
              for (k=j+1; k<=i; k++){
                virt_blk_len_A[k] = 1;
                new_sym_A[k] = NS;
              }
              break;
            }
          }
          i_A++;
        }
      }
solomon's avatar
solomon committed
709
      itsr = B->rec_tsr;
710
      i_B = 0;
solomon's avatar
solomon committed
711 712
      for (i=0; i<B->order; i++){
        if (B->sym[i] == NS){
713
          for (j=0; j<itsr->order; j++){
solomon's avatar
solomon committed
714
            if (B->inner_ordering[j] == i_B){
715 716 717
              j=i;
              do {
                j--;
solomon's avatar
solomon committed
718
              } while (j>=0 && B->sym[j] != NS);
719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735
              for (k=j+1; k<=i; k++){
                virt_blk_len_B[k] = 1;
                new_sym_B[k] = NS;
              }
              break;
            }
          }
          i_B++;
        }
      }
    }
    if (is_top) {
      htsum = tsumseq;
      is_top = 0;
    } else {
      *rec_tsum = tsumseq;
    }
736 737
//    tsumseq->order_A     = A->order;
//    tsumseq->idx_map_A   = idx_A;
738 739
    tsumseq->edge_len_A  = virt_blk_len_A;
    tsumseq->sym_A       = new_sym_A;
740 741
//    tsumseq->order_B     = B->order;
//    tsumseq->idx_map_B   = idx_B;
742 743 744
    tsumseq->edge_len_B  = virt_blk_len_B;
    tsumseq->sym_B       = new_sym_B;
    tsumseq->is_custom   = is_custom;
solomon's avatar
solomon committed
745
    if (is_custom){
746 747 748
      tsumseq->is_inner  = 0;
      tsumseq->func      = func;
    } else tsumseq->func = NULL;
749 750
//    htsum->alpha         = alpha;
//    htsum->beta          = beta;
751

752 753
//    htsum->A = A->data;
//    htsum->B = B->data;
754

solomon's avatar
solomon committed
755 756 757 758
    CTF_int::cdealloc(idx_arr);
    CTF_int::cdealloc(blk_len_A);
    CTF_int::cdealloc(blk_len_B);
    CTF_int::cdealloc(phys_mapped);
759 760 761 762 763 764

    return htsum;
  }

  int summation::home_sum_tsr(bool run_diag){
    int ret, was_home_A, was_home_B;
solomon's avatar
solomon committed
765
    tensor * tnsr_A, * tnsr_B;
766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790
    // code below turns summations into scaling, but never seems to be invoked in AQ or test_suite, so commenting it out for now
/*    if (A==B && !is_custom){
      bool is_scal = true;
      for (int i=0; i<A->order; i++){
        if (idx_A[i] != idx_B[i]) is_scal = false;
      }
      if (is_scal){
        if (alpha == NULL && beta == NULL){
          scaling scl = scaling(A, idx_A, NULL);
          scl.execute();
        } else {
          char nalpha[A->sr->el_size];
          if (alpha == NULL) A->sr->copy(nalpha, A->sr->mulid());
          else A->sr->copy(nalpha, alpha);

          if (beta == NULL) A->sr->add(nalpha, A->sr->mulid(), nalpha);
          else A->sr->add(nalpha, beta, nalpha);

          scaling scl = scaling(A, idx_A, nalpha);
          scl.execute();
        }
        return SUCCESS;
      }
    }*/

solomon's avatar
solomon committed
791 792
    summation osum = summation(*this);
   
793
    CTF_int::contract_mst();
794

795 796
    A->unfold();
    B->unfold();
797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813
    // FIXME: if custom function, we currently don't know whether its odd, even or neither, so unpack everything
    if (is_custom){
      bool is_nonsym=true;
      for (int i=0; i<A->order; i++){
        if (A->sym[i] != NS){
          is_nonsym = false;
        }
      }
      if (!is_nonsym){
        int sym_A[A->order];
        std::fill(sym_A, sym_A+A->order, NS);
        int idx_A[A->order];
        for (int i=0; i<A->order; i++){
          idx_A[i] = i;
        }
        tensor tA(A->sr, A->order, A->lens, sym_A, A->wrld, 1);
        tA.is_home = 0;
814
        tA.has_home = 0;
815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851
        summation st(A, idx_A, A->sr->mulid(), &tA, idx_A, A->sr->mulid());
        st.execute();
        summation stme(*this);
        stme.A = &tA;
        stme.execute();
        return SUCCESS;
      }
    }
    if (is_custom){
      bool is_nonsym=true;
      for (int i=0; i<B->order; i++){
        if (B->sym[i] != NS){
          is_nonsym = false;
        }
      }
      if (!is_nonsym){
        int sym_B[B->order];
        std::fill(sym_B, sym_B+B->order, NS);
        int idx_B[B->order];
        for (int i=0; i<B->order; i++){
          idx_B[i] = i;
        }
        tensor tB(B->sr, B->order, B->lens, sym_B, B->wrld, 1);
        tB.is_home = 0;
        if (!B->sr->isequal(B->sr->addid(), beta)){
          summation st(B, idx_B, B->sr->mulid(), &tB, idx_B, B->sr->mulid());
          st.execute();
        }
        summation stme(*this);
        stme.B = &tB;
        stme.execute();
        summation stme2(&tB, idx_B, B->sr->mulid(), B, idx_B, B->sr->addid());
        stme2.execute();
        return SUCCESS;
      }
    }

852 853
  #ifndef HOME_CONTRACT
    #ifdef USE_SYM_SUM
solomon's avatar
solomon committed
854
      ret = sym_sum_tsr(run_diag);
855 856
      return ret;
    #else
solomon's avatar
solomon committed
857
      ret = sum_tensors(run_diag);
858 859 860
      return ret;
    #endif
  #else
solomon's avatar
solomon committed
861 862
    if (A->has_zero_edge_len || 
        B->has_zero_edge_len){
863
      if (!B->sr->isequal(beta,B->sr->mulid()) && !B->has_zero_edge_len){ 
solomon's avatar
solomon committed
864
        int sub_idx_map_B[B->order];
865
        int sm_idx=0;
solomon's avatar
solomon committed
866
        for (int i=0; i<B->order; i++){
867 868 869
          sub_idx_map_B[i]=sm_idx;
          sm_idx++;
          for (int j=0; j<i; j++){
solomon's avatar
solomon committed
870
            if (idx_B[i]==idx_B[j]){
871 872 873 874 875 876
              sub_idx_map_B[i]=sub_idx_map_B[j];
              sm_idx--;
              break;
            }
          }
        }
solomon's avatar
solomon committed
877 878
        scaling scl = scaling(B, sub_idx_map_B, beta);
        scl.execute();
879
      }
solomon's avatar
solomon committed
880
      return SUCCESS;
881
    }
solomon's avatar
solomon committed
882 883 884 885
    if (A == B){
      tensor * cpy_tsr_A = new tensor(A);
      osum.A = cpy_tsr_A;
      osum.execute();
886
      delete cpy_tsr_A;
solomon's avatar
solomon committed
887
      return SUCCESS;
888
    }
solomon's avatar
solomon committed
889 890
    was_home_A = A->is_home;
    was_home_B = B->is_home;
891
    if (was_home_A){
892 893
      tnsr_A              = new tensor(A,0,0);
      tnsr_A->data        = A->data;
solomon's avatar
solomon committed
894
      tnsr_A->home_buffer = A->home_buffer;
895 896 897
      tnsr_A->is_home     = 1;
      tnsr_A->is_mapped   = 1;
      tnsr_A->topo        = A->topo;
solomon's avatar
solomon committed
898 899
      copy_mapping(A->order, A->edge_map, tnsr_A->edge_map);
      tnsr_A->set_padding();
900
      osum.A              = tnsr_A;
901
    } else tnsr_A = NULL;     
902
    if (was_home_B){
903 904
      tnsr_B              = new tensor(B,0,0);
      tnsr_B->data        = B->data;
solomon's avatar
solomon committed
905
      tnsr_B->home_buffer = B->home_buffer;
906 907 908
      tnsr_B->is_home     = 1;
      tnsr_B->is_mapped   = 1;
      tnsr_B->topo        = B->topo;
solomon's avatar
solomon committed
909 910
      copy_mapping(B->order, B->edge_map, tnsr_B->edge_map);
      tnsr_B->set_padding();
911
      osum.B              = tnsr_B;
912
    } else tnsr_B = NULL;
913
  #if DEBUG >= 1
solomon's avatar
solomon committed
914
    if (A->wrld->cdt.rank == 0)
915 916 917 918
      printf("Start head sum:\n");
  #endif
    
    #ifdef USE_SYM_SUM
solomon's avatar
solomon committed
919
    ret = osum.sym_sum_tsr(run_diag);
920
    #else
solomon's avatar
solomon committed
921
    ret = osum.sum_tensors(run_diag);
922 923
    #endif
  #if DEBUG >= 1
solomon's avatar
solomon committed
924
    if (A->wrld->cdt.rank == 0)
925 926 927
      printf("End head sum:\n");
  #endif

solomon's avatar
solomon committed
928
    if (ret!= SUCCESS) return ret;
929 930
    if (was_home_A) tnsr_A->unfold(); 
    else A->unfold();
solomon's avatar
solomon committed
931
    if (was_home_B) tnsr_B->unfold();
932
    else B->unfold();
933

solomon's avatar
solomon committed
934 935
    if (was_home_B && !tnsr_B->is_home){
      if (A->wrld->cdt.rank == 0)
936
        DPRINTF(2,"Migrating tensor %s back to home\n", B->name);
937
      distribution odst(tnsr_B);
solomon's avatar
solomon committed
938 939
      B->data = tnsr_B->data;
      B->is_home = 0;
940
      TAU_FSTART(redistribute_for_sum_home);
solomon's avatar
solomon committed
941
      B->redistribute(odst);
942
      TAU_FSTOP(redistribute_for_sum_home);
943
      memcpy(B->home_buffer, B->data, B->size*B->sr->el_size);
solomon's avatar
solomon committed
944
      CTF_int::cdealloc(B->data);
solomon's avatar
solomon committed
945 946 947 948
      B->data = B->home_buffer;
      B->is_home = 1;
      tnsr_B->is_data_aliased = 1;
      delete tnsr_B;
949
    } else if (was_home_B){
solomon's avatar
solomon committed
950 951
      if (tnsr_B->data != B->data){
        printf("Tensor %s is a copy of %s and did not leave home but buffer is %p was %p\n", tnsr_B->name, B->name, tnsr_B->data, B->data);
952 953 954
        ABORT;

      }
solomon's avatar
solomon committed
955 956 957
      tnsr_B->has_home = 0;
      tnsr_B->is_data_aliased = 1;
      delete tnsr_B;
958
    }
solomon's avatar
solomon committed
959 960 961
    if (was_home_A && !tnsr_A->is_home){
      tnsr_A->has_home = 0;
      delete tnsr_A;
962
    } else if (was_home_A) {
solomon's avatar
solomon committed
963 964 965
      tnsr_A->has_home = 0;
      tnsr_A->is_data_aliased = 1;
      delete tnsr_A;
966 967 968 969 970 971
    }
    return ret;
  #endif
  }

  int summation::sym_sum_tsr(bool run_diag){
972
    int sidx, i, nst_B, * new_idx_map;
solomon's avatar
solomon committed
973
    int * map_A, * map_B;
974
    int ** dstack_map_B;
975
    tensor * tnsr_A, * tnsr_B, * new_tsr, ** dstack_tsr_B;
solomon's avatar
solomon committed
976
    std::vector<summation> perm_types;
977
    std::vector<int> signs;
solomon's avatar
solomon committed
978
    char const * dbeta;
979 980 981
  #if (DEBUG >= 2 || VERBOSE >= 1)
    print();
  #endif
982
    check_consistency();
983 984 985

    A->unfold();
    B->unfold();
986
    if (A->has_zero_edge_len || B->has_zero_edge_len){
987
      if (!B->sr->isequal(beta, B->sr->mulid()) && !B->has_zero_edge_len){ 
988
        int sub_idx_map_B[B->order];
989
        int sm_idx=0;
990
        for (int i=0; i<B->order; i++){
991 992 993
          sub_idx_map_B[i]=sm_idx;
          sm_idx++;
          for (int j=0; j<i; j++){
994
            if (idx_B[i]==idx_B[j]){
995 996 997 998 999 1000
              sub_idx_map_B[i]=sub_idx_map_B[j];
              sm_idx--;
              break;
            }
          }
        }
1001 1002
        scaling scl = scaling(B, sub_idx_map_B, beta);
        scl.execute();
1003
      }
solomon's avatar
solomon committed
1004
      return SUCCESS;
1005
    }
1006
    // If we have sparisity, use separate mechanism
1007
    /*if (A->is_sparse || B->is_sparse){
1008 1009
      sp_sum();
      return SUCCESS;
1010
    }*/
1011 1012
    tnsr_A = A;
    tnsr_B = B;
1013
    char * new_alpha = (char*)alloc(tnsr_B->sr->el_size);
1014 1015 1016 1017
    CTF_int::alloc_ptr(sizeof(int)*tnsr_A->order,     (void**)&map_A);
    CTF_int::alloc_ptr(sizeof(int)*tnsr_B->order,     (void**)&map_B);
    CTF_int::alloc_ptr(sizeof(int*)*tnsr_B->order,    (void**)&dstack_map_B);
    CTF_int::alloc_ptr(sizeof(tensor*)*tnsr_B->order, (void**)&dstack_tsr_B);
1018 1019 1020 1021
    memcpy(map_A, idx_A, tnsr_A->order*sizeof(int));
    memcpy(map_B, idx_B, tnsr_B->order*sizeof(int));
    while (!run_diag && tnsr_A->extract_diag(map_A, 1, new_tsr, &new_idx_map) == SUCCESS){
      if (tnsr_A != A) delete tnsr_A;
solomon's avatar
solomon committed
1022
      CTF_int::cdealloc(map_A);
1023
      tnsr_A = new_tsr;
1024 1025 1026
      map_A = new_idx_map;
    }
    nst_B = 0;
1027
    while (!run_diag && tnsr_B->extract_diag(map_B, 1, new_tsr, &new_idx_map) == SUCCESS){
1028
      dstack_map_B[nst_B] = map_B;
1029
      dstack_tsr_B[nst_B] = tnsr_B;
1030
      nst_B++;
1031
      tnsr_B = new_tsr;
1032 1033 1034
      map_B = new_idx_map;
    }

Edgar Solomonik's avatar
Edgar Solomonik committed
1035 1036 1037
    summation new_sum = summation(*this);
    new_sum.A = tnsr_A;
    new_sum.B = tnsr_B;
1038 1039
    memcpy(new_sum.idx_A, map_A, sizeof(int)*tnsr_A->order);
    memcpy(new_sum.idx_B, map_B, sizeof(int)*tnsr_B->order);
1040
    if (tnsr_A == tnsr_B){
1041
      tensor nnew_tsr = tensor(tnsr_A);
Edgar Solomonik's avatar
Edgar Solomonik committed
1042 1043
      new_sum.A = &nnew_tsr;
      new_sum.B = tnsr_B;
1044
      new_sum.sym_sum_tsr(run_diag);
1045 1046
      
      /*clone_tensor(ntid_A, 1, &new_tid);
1047 1048 1049 1050
      new_type = *type;
      new_type.tid_A = new_tid;
      stat = sym_sum_tsr(alpha_, beta, &new_type, ftsr, felm, run_diag);
      del_tsr(new_tid);
1051
      return stat;*/
1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066
    } else {
      
  /*    new_type.tid_A = ntid_A;
      new_type.tid_B = ntid_B;
      new_type.idx_map_A = map_A;
      new_type.idx_map_B = map_B;*/
  
      //FIXME: make these typefree...
      int sign = align_symmetric_indices(tnsr_A->order,
                                         new_sum.idx_A,
                                         tnsr_A->sym,
                                         tnsr_B->order,
                                         new_sum.idx_B,
                                         tnsr_B->sym);
      int ocfact = overcounting_factor(tnsr_A->order,
1067
                                       new_sum.idx_A,
1068 1069
                                       tnsr_A->sym,
                                       tnsr_B->order,
1070
                                       new_sum.idx_B,
1071
                                       tnsr_B->sym);
1072 1073 1074
  
      if (ocfact != 1 || sign != 1){
        if (ocfact != 1){
1075
          tnsr_B->sr->safecopy(new_alpha, tnsr_B->sr->addid());
1076 1077 1078 1079 1080 1081 1082
          
          for (int i=0; i<ocfact; i++){
            tnsr_B->sr->add(new_alpha, alpha, new_alpha);
          }
          alpha = new_alpha;
        }
        if (sign == -1){
1083
          tnsr_B->sr->safeaddinv(alpha, new_alpha);
1084
          alpha = new_alpha;
1085 1086
        }
      }
1087
  
1088
      /*if (A->sym[0] == SY && B->sym[0] == AS){
solomon's avatar
solomon committed
1089 1090
        print();
        ASSERT(0); 
1091
      }*/
1092
      if (new_sum.unfold_broken_sym(NULL) != -1){
solomon's avatar
solomon committed
1093
        if (A->wrld->cdt.rank == 0)
1094
          DPRINTF(1, "Permutational symmetry is broken\n");
1095 1096 1097 1098 1099
  
        summation * unfold_sum;
        sidx = new_sum.unfold_broken_sym(&unfold_sum);
        int sy;
        sy = 0;
1100
        int sidx2 = unfold_sum->unfold_broken_sym(NULL);
1101 1102
        if (sidx%2 == 0 && (A->sym[sidx/2] == SY || unfold_sum->A->sym[sidx/2] == SY)) sy = 1;
        if (sidx%2 == 1 && (B->sym[sidx/2] == SY || unfold_sum->B->sym[sidx/2] == SY)) sy = 1;
1103
        //if (sy && sidx%2 == 0){
1104 1105 1106 1107 1108
        if (sidx2 != -1 || 
            (sy && (sidx%2 == 0  || !tnsr_B->sr->isequal(new_sum.beta, tnsr_B->sr->addid())))){
          if (sidx%2 == 0){
            if (unfold_sum->A->sym[sidx/2] == NS){
              if (A->wrld->cdt.rank == 0)
1109
                DPRINTF(1,"Performing operand desymmetrization for summation of A idx=%d\n",sidx/2);
1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147
              desymmetrize(tnsr_A, unfold_sum->A, 0);
            } else {
              if (A->wrld->cdt.rank == 0)
                DPRINTF(1,"Performing operand symmetrization for summation\n");
              symmetrize(unfold_sum->A, tnsr_A);
            }
            //unfold_sum->B = tnsr_B;
            unfold_sum->sym_sum_tsr(run_diag);
    //        sym_sum_tsr(alpha, beta, &unfold_type, ftsr, felm, run_diag);
            if (tnsr_A != unfold_sum->A){
              unfold_sum->A->unfold();
              tnsr_A->pull_alias(unfold_sum->A);
              delete unfold_sum->A;
            }
          } else {
            //unfold_sum->A = tnsr_A;
            if (A->wrld->cdt.rank == 0)
              DPRINTF(1,"Performing product desymmetrization for summation\n");
            desymmetrize(tnsr_B, unfold_sum->B, 1);
            unfold_sum->sym_sum_tsr(run_diag);
            if (A->wrld->cdt.rank == 0)
              DPRINTF(1,"Performing product symmetrization for summation\n");
            if (tnsr_B->data != unfold_sum->B->data && !tnsr_B->sr->isequal(tnsr_B->sr->mulid(), unfold_sum->beta)){
              int sidx_B[tnsr_B->order];
              for (int iis=0; iis<tnsr_B->order; iis++){
                sidx_B[iis] = iis;
              }
              scaling sscl = scaling(tnsr_B, sidx_B, unfold_sum->beta);
              sscl.execute();
            }
            symmetrize(tnsr_B, unfold_sum->B);

    //        sym_sum_tsr(alpha, beta, &unfold_type, ftsr, felm, run_diag);
            if (tnsr_B != unfold_sum->B){
              unfold_sum->B->unfold();
              tnsr_B->pull_alias(unfold_sum->B);
              delete unfold_sum->B;
            }
1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161
          }
        } else {
          if (sidx != -1 && sidx%2 == 1){
            delete unfold_sum->B;
          } else if (sidx != -1 && sidx%2 == 0){
            delete unfold_sum->A;
          }
          //get_sym_perms(&new_type, alpha, perm_types, signs);
          get_sym_perms(new_sum, perm_types, signs);
          if (A->wrld->cdt.rank == 0)
            DPRINTF(1,"Performing %d summation permutations\n",
                    (int)perm_types.size());
          dbeta = beta;
          char * new_alpha = (char*)alloc(tnsr_B->sr->el_size);
1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174

          tensor * inv_tsr_A;
          bool need_inv = false;
          // if we have no multiplicative operator, must inverse sign manually
          if (tnsr_B->sr->mulid() == NULL){
            for (i=0; i<(int)perm_types.size(); i++){
              need_inv = true;
            }
            if (need_inv){
              inv_tsr_A = new tensor(tnsr_A);
              inv_tsr_A->addinv();
            }
          }
1175
          for (i=0; i<(int)perm_types.size(); i++){
1176 1177 1178 1179 1180 1181 1182 1183 1184 1185
            // if group apply additive inverse manually
            if (signs[i] == -1 && need_inv){
              perm_types[i].A = inv_tsr_A;
            } else {
              if (signs[i] == 1)
                tnsr_B->sr->safecopy(new_alpha, alpha);
              else 
                tnsr_B->sr->safeaddinv(alpha, new_alpha);
              perm_types[i].alpha = new_alpha;
            }
1186 1187 1188 1189
            perm_types[i].beta = dbeta;
            perm_types[i].sum_tensors(run_diag);
            dbeta = new_sum.B->sr->mulid();
          }
solomon's avatar
solomon committed
1190
          cdealloc(new_alpha);
1191 1192
          if (need_inv)
            delete inv_tsr_A;
1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204
  /*        for (i=0; i<(int)perm_types.size(); i++){
            free_type(&perm_types[i]);
          }*/
          perm_types.clear();
          signs.clear();
        }
        delete unfold_sum;
      } else {
        new_sum.alpha = alpha;
        new_sum.sum_tensors(run_diag);
  /*      sum_tensors(alpha, beta, new_type.tid_A, new_type.tid_B, new_type.idx_map_A,
                    new_type.idx_map_B, ftsr, felm, run_diag);*/
1205 1206
      }
    }
1207
    if (tnsr_A != A) delete tnsr_A;
1208
    for (i=nst_B-1; i>=0; i--){
1209 1210 1211 1212 1213
//      extract_diag(dstack_tid_B[i], dstack_map_B[i], 0, &ntid_B, &new_idx_map);
      dstack_tsr_B[i]->extract_diag(dstack_map_B[i], 0, tnsr_B, &new_idx_map);
      //del_tsr(ntid_B);
      delete tnsr_B;
      tnsr_B = dstack_tsr_B[i];
1214
    }
1215
    ASSERT(tnsr_B == B);
solomon's avatar
solomon committed
1216 1217 1218 1219 1220
    CTF_int::cdealloc(new_alpha);
    CTF_int::cdealloc(map_A);
    CTF_int::cdealloc(map_B);
    CTF_int::cdealloc(dstack_map_B);
    CTF_int::cdealloc(dstack_tsr_B);
1221

solomon's avatar
solomon committed
1222
    return SUCCESS;
1223 1224 1225 1226 1227 1228 1229 1230
  }


  /**
   * \brief PDAXPY: a*idx_map_A(A) + b*idx_map_B(B) -> idx_map_B(B).
   * \param[in] run_diag if 1 run diagonal sum
   */
  int summation::sum_tensors(bool run_diag){
1231
    int stat, * new_idx_map;
1232
    int * map_A, * map_B;
1233
    int nst_B;
1234
    int ** dstack_map_B;
1235 1236
    tensor * tnsr_A, * tnsr_B, * new_tsr, ** dstack_tsr_B;
//    tsum<dtype> * sumf;
solomon's avatar
solomon committed
1237
    tsum * sumf;
1238 1239 1240
    //check_sum(tid_A, tid_B, idx_map_A, idx_map_B);
    //FIXME: hmm all of the below already takes place in sym_sum
    check_consistency();
1241 1242
    A->unfold();
    B->unfold();
1243
    if (A->has_zero_edge_len || B->has_zero_edge_len){
1244
      if (!B->sr->isequal(beta,B->sr->mulid()) && !B->has_zero_edge_len){ 
1245
    /*    fseq_scl<dtype> fs;
1246 1247
        fs.func_ptr=sym_seq_scl_ref<dtype>;
        fseq_elm_scl<dtype> felm;
1248
        felm.func_ptr = NULL;*/
1249
        int sub_idx_map_B[B->order];
1250
        int sm_idx=0;
1251
        for (int i=0; i<B->order; i++){
1252 1253 1254
          sub_idx_map_B[i]=sm_idx;
          sm_idx++;
          for (int j=0; j<i; j++){
1255
            if (idx_B[i]==idx_B[j]){
1256 1257 1258 1259 1260 1261
              sub_idx_map_B[i]=sub_idx_map_B[j];
              sm_idx--;
              break;
            }
          }
        }
1262 1263
        scaling scl = scaling(B, sub_idx_map_B, beta);
        scl.execute();
1264
      }
solomon's avatar
solomon committed
1265
      return SUCCESS;
1266 1267 1268
    }


1269
    //FIXME: remove all of the below, sum_tensors should never be called without sym_sum
1270 1271 1272 1273
    CTF_int::alloc_ptr(sizeof(int)*A->order,  (void**)&map_A);
    CTF_int::alloc_ptr(sizeof(int)*B->order,  (void**)&map_B);
    CTF_int::alloc_ptr(sizeof(int*)*B->order, (void**)&dstack_map_B);
    CTF_int::alloc_ptr(sizeof(tensor*)*B->order, (void**)&dstack_tsr_B);
1274 1275 1276 1277 1278 1279
    tnsr_A = A;
    tnsr_B = B;
    memcpy(map_A, idx_A, tnsr_A->order*sizeof(int));
    memcpy(map_B, idx_B, tnsr_B->order*sizeof(int));
    while (!run_diag && tnsr_A->extract_diag(map_A, 1, new_tsr, &new_idx_map) == SUCCESS){
      if (tnsr_A != A) delete tnsr_A;
solomon's avatar
solomon committed
1280
      CTF_int::cdealloc(map_A);
1281
      tnsr_A = new_tsr;
1282 1283 1284
      map_A = new_idx_map;
    }
    nst_B = 0;
1285
    while (!run_diag && tnsr_B->extract_diag(map_B, 1, new_tsr, &new_idx_map) == SUCCESS){
1286
      dstack_map_B[nst_B] = map_B;
1287
      dstack_tsr_B[nst_B] = tnsr_B;
1288
      nst_B++;
1289
      tnsr_B = new_tsr;
1290 1291
      map_B = new_idx_map;
    }
1292 1293 1294 1295 1296 1297 1298 1299 1300

    if (!tnsr_A->is_sparse && tnsr_B->is_sparse){
      tensor * stnsr_A = tnsr_A;
      tnsr_A = new tensor(stnsr_A);
      tnsr_A->sparsify(); 
      if (A != stnsr_A) delete stnsr_A;

    }

1301
    // FIXME: if A has self indices and function is distributive, presum A first, otherwise if it is dense and B is sparse, sparsify A
Edgar Solomonik's avatar
Edgar Solomonik committed
1302 1303 1304
    summation new_sum = summation(*this);
    new_sum.A = tnsr_A;
    new_sum.B = tnsr_B;
1305
    if (tnsr_A == tnsr_B){
solomon's avatar
solomon committed
1306
      tensor * nnew_tsr = new tensor(tnsr_A);
Edgar Solomonik's avatar
Edgar Solomonik committed
1307 1308
      new_sum.A = nnew_tsr;
      new_sum.B = tnsr_B;
1309
    } else{ 
1310 1311
     //FIXME: remove the below, sum_tensors should never be called without sym_sum
     int sign = align_symmetric_indices(tnsr_A->order,
1312
                                        new_sum.idx_A,
1313 1314
                                        tnsr_A->sym,
                                        tnsr_B->order,
1315
                                        new_sum.idx_B,
1316
                                        tnsr_B->sym);
1317 1318 1319 1320 1321

      #if DEBUG >= 1
        new_sum.print();
      #endif

1322
      ASSERT(sign == 1);
solomon's avatar
solomon committed
1323
/*        if (sign == -1){
1324 1325
          char * new_alpha = (char*)malloc(tnsr_B->sr->el_size);
          tnsr_B->sr->addinv(alpha, new_alpha);
solomon's avatar
solomon committed
1326 1327
          alpha = new_alpha;
        }*/
1328

1329
  #if 0 //VERIFY
1330 1331 1332 1333 1334 1335 1336 1337
      int64_t nsA, nsB;
      int64_t nA, nB;
      dtype * sA, * sB;
      dtype * uA, * uB;
      int order_A, order_B,  i;
      int * edge_len_A, * edge_len_B;
      int * sym_A, * sym_B;
      stat = allread_tsr(ntid_A, &nsA, &sA);
solomon's avatar
solomon committed
1338
      assert(stat == SUCCESS);
1339 1340

      stat = allread_tsr(ntid_B, &nsB, &sB);
solomon's avatar
solomon committed
1341
      assert(stat == SUCCESS);
1342 1343 1344 1345 1346 1347 1348 1349
  #endif

      TAU_FSTART(sum_tensors);

      /* Check if the current tensor mappings can be summed on */
  #if REDIST
      if (1) {
  #else
1350
      if (new_sum.check_mapping() == 0) {
1351 1352
  #endif
        /* remap if necessary */
1353
        stat = new_sum.map();
1354
        if (stat == ERROR) {
1355
          printf("Failed to map tensors to physical grid\n");
1356
          return ERROR;
1357 1358
        }
      } else {
1359
  #if DEBUG > 2
solomon's avatar
solomon committed
1360
        if (A->wrld->cdt.rank == 0){
1361 1362
          printf("Keeping mappings:\n");
        }
1363 1364
        tnsr_A->print_map(stdout);
        tnsr_B->print_map(stdout);
1365 1366 1367
  #endif
      }
      /* Construct the tensor algorithm we would like to use */
1368
      ASSERT(new_sum.check_mapping());
1369
  #if FOLD_TSR
1370
      if (is_custom == false && new_sum.can_fold()){
1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386
        //FIXME bit of a guess, no?
        double est_time_nofold = 4.*(A->size + B->size)*COST_MEMBW;
        if (new_sum.est_time_fold() + (A->size + B->size)*COST_MEMBW < est_time_nofold){
          /*if (A->wrld->cdt.rank == 0)
            printf("Decided to fold\n");*/
          int inner_stride;
          TAU_FSTART(map_fold);
          inner_stride = new_sum.map_fold();
          TAU_FSTOP(map_fold);
          sumf = new_sum.construct_sum(inner_stride);
        } else {
          /*if (A->wrld->cdt.rank == 0)
            printf("Decided not to fold\n");*/
        
          sumf = new_sum.construct_sum();
        }
1387 1388 1389 1390 1391 1392
      } else{
  #if DEBUG >= 1
        if (A->wrld->cdt.rank == 0){
          printf("Could not fold summation, is_custom = %d, new_sum.can_fold = %d\n", is_custom, new_sum.can_fold());
        }
  #endif
Edgar Solomonik's avatar
Edgar Solomonik committed
1393
        sumf = new_sum.construct_sum();
1394
      }
solomon's avatar
solomon committed
1395 1396
        /*sumf = construct_sum(alpha, beta, ntid_A, map_A, ntid_B, map_B,
                             ftsr, felm);*/
1397
  #else
Edgar Solomonik's avatar
Edgar Solomonik committed
1398
      sumf = new_sum.construct_sum();
solomon's avatar
solomon committed
1399 1400
      /*sumf = construct_sum(alpha, beta, ntid_A, map_A, ntid_B, map_B,
                           ftsr, felm);*/
1401 1402 1403 1404 1405 1406 1407
  #endif
      /*TAU_FSTART(zero_sum_padding);
      stat = zero_out_padding(ntid_A);
      TAU_FSTOP(zero_sum_padding);
      TAU_FSTART(zero_sum_padding);
      stat = zero_out_padding(ntid_B);
      TAU_FSTOP(zero_sum_padding);*/
solomon's avatar
solomon committed
1408
      DEBUG_PRINTF("[%d] performing tensor sum\n", A->wrld->cdt.rank);
1409
  #if DEBUG >=3
1410
      /*if (A->wrld->cdt.rank == 0){
1411 1412 1413 1414 1415 1416
        for (int i=0; i<tensors[ntid_A]->order; i++){
          printf("padding[%d] = %d\n",i, tensors[ntid_A]->padding[i]);
        }
        for (int i=0; i<tensors[ntid_B]->order; i++){
          printf("padding[%d] = %d\n",i, tensors[ntid_B]->padding[i]);
        }
1417
      }*/
1418 1419
  #endif

1420 1421 1422
  #if DEBUG >= 2
      if (tnsr_B->wrld->rank==0)
        sumf->print();
1423 1424
      tnsr_A->print_map();
      tnsr_B->print_map();
1425
  #endif
1426 1427
      TAU_FSTART(sum_func);
      /* Invoke the contraction algorithm */
1428
      tnsr_A->topo->activate();
1429
      MPI_Barrier(tnsr_B->wrld->comm);
1430
      sumf->run();
1431 1432 1433 1434 1435 1436 1437 1438 1439
      if (tnsr_B->is_sparse){
        if (tnsr_B->data != sumf->new_B){
          cdealloc(tnsr_B->data);
          tnsr_B->data = sumf->new_B;
          //tnsr_B->nnz_loc = sumf->new_nnz_B;
          tnsr_B->nnz_loc = 0;
          for (int i=0; i<tnsr_B->calc_nvirt(); i++){
            tnsr_B->nnz_loc += tnsr_B->nnz_blk[i];
          }
1440
        }
1441
        printf("tnsr_B->data = %p, nnz_loc = %ld\n", tnsr_B->data, tnsr_B->nnz_loc);
1442
      }
1443
      /*tnsr_B->unfold();
1444 1445 1446 1447 1448 1449 1450 1451
      tnsr_B->print();
      MPI_Barrier(tnsr_B->wrld->comm);
      if (tnsr_B->wrld->rank==1){
      for (int i=0; i<tnsr_B->size; i++){
        printf("[%d] %dth element  ",tnsr_B->wrld->rank,i);
        tnsr_B->sr->print(tnsr_B->data+i*tnsr_B->sr->el_size);
        printf("\n");
      }
1452 1453
      }*/
      tnsr_A->topo->deactivate();
solomon's avatar
solomon committed
1454 1455
      tnsr_A->unfold();
      tnsr_B->unfold();
1456 1457 1458
#ifndef SEQ
      stat = tnsr_B->zero_out_padding();
#endif
1459
      TAU_FSTOP(sum_func);
1460

1461
  #if 0 //VERIFY
1462
      stat = allread_tsr(ntid_A, &nA, &uA);
solomon's avatar
solomon committed
1463
      assert(stat == SUCCESS);
1464
      stat = get_info(ntid_A, &order_A, &edge_len_A, &sym_A);
solomon's avatar
solomon committed
1465
      assert(stat == SUCCESS);
1466 1467

      stat = allread_tsr(ntid_B, &nB, &uB);
solomon's avatar
solomon committed
1468
      assert(stat == SUCCESS);
1469
      stat = get_info(ntid_B, &order_B, &edge_len_B, &sym_B);
solomon's avatar
solomon committed
1470
      assert(stat == SUCCESS);
1471 1472 1473

      if (nsA != nA) { printf("nsA = " PRId64 ", nA = " PRId64 "\n",nsA,nA); ABORT; }
      if (nsB != nB) { printf("nsB = " PRId64 ", nB = " PRId64 "\n",nsB,nB); ABORT; }
1474
      for (i=0; (int64_t)i<nA; i++){
1475 1476 1477 1478 1479 1480 1481
        if (fabs(uA[i] - sA[i]) > 1.E-6){
          printf("A[i] = %lf, sA[i] = %lf\n", uA[i], sA[i]);
        }
      }

      cpy_sym_sum(alpha, uA, order_A, edge_len_A, edge_len_A, sym_A, map_A,
                  beta, sB, order_B, edge_len_B, edge_len_B, sym_B, map_B);
solomon's avatar
solomon committed
1482
      assert(stat == SUCCESS);
1483

1484
      for (i=0; (int64_t)i<nB; i++){
1485 1486 1487 1488 1489
        if (fabs(uB[i] - sB[i]) > 1.E-6){
          printf("B[%d] = %lf, sB[%d] = %lf\n", i, uB[i], i, sB[i]);
        }
        assert(fabs(sB[i] - uB[i]) < 1.E-6);
      }
solomon's avatar
solomon committed
1490 1491 1492 1493
      CTF_int::cdealloc(uA);
      CTF_int::cdealloc(uB);
      CTF_int::cdealloc(sA);
      CTF_int::cdealloc(sB);
1494 1495 1496
  #endif

      delete sumf;
solomon's avatar
solomon committed
1497
      if (tnsr_A != A) delete tnsr_A;
1498
      for (int i=nst_B-1; i>=0; i--){
solomon's avatar
solomon committed
1499
        int ret = dstack_tsr_B[i]->extract_diag(dstack_map_B[i], 0, tnsr_B, &new_idx_map);
solomon's avatar
solomon committed
1500
        ASSERT(ret == SUCCESS);
solomon's avatar
solomon committed
1501 1502
        delete tnsr_B;
        tnsr_B = dstack_tsr_B[i];
1503
      }
solomon's avatar
solomon committed
1504
      ASSERT(tnsr_B == B);
1505
    }
1506
  //#ifndef SEQ
1507
    //stat = B->zero_out_padding();
1508
  //#endif
solomon's avatar
solomon committed
1509 1510 1511 1512
    CTF_int::cdealloc(map_A);
    CTF_int::cdealloc(map_B);
    CTF_int::cdealloc(dstack_map_B);
    CTF_int::cdealloc(dstack_tsr_B);
1513 1514

    TAU_FSTOP(sum_tensors);
solomon's avatar
solomon committed
1515
    return SUCCESS;
1516 1517
  }

1518 1519 1520
  int summation::unfold_broken_sym(summation ** nnew_sum){
    int sidx, i, num_tot, iA, iA2, iB;
    int * idx_arr;
1521
    int bsym = NS;
solomon's avatar
solomon committed
1522
    summation * new_sum;
1523 1524 1525 1526
   
    if (nnew_sum != NULL){
      new_sum = new summation(*this);
      *nnew_sum = new_sum;
1527
    } else new_sum = NULL;
1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538

    inv_idx(A->order, idx_A,
            B->order, idx_B,
            &num_tot, &idx_arr);

    sidx = -1;
    for (i=0; i<A->order; i++){
      if (A->sym[i] != NS){
        iA = idx_A[i];
        if (idx_arr[2*iA+1] != -1){
          if (B->sym[idx_arr[2*iA+1]] == NS ||
1539
              ((B->sym[idx_arr[2*iA+1]] == AS) ^ (A->sym[i] == AS)) ||
1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556
              idx_arr[2*idx_A[i+1]+1] == -1 ||
              idx_A[i+1] != idx_B[idx_arr[2*iA+1]+1]){
            sidx = 2*i;
            break;
          }
        } else if (idx_arr[2*idx_A[i+1]+1] != -1){
          sidx = 2*i;
          break;
        }
      }
    } 
    if (sidx == -1){
      for (i=0; i<B->order; i++){
        if (B->sym[i] != NS){
          iB = idx_B[i];
          if (idx_arr[2*iB+0] != -1){
            if (A->sym[idx_arr[2*iB+0]] == NS ||
1557
                ((A->sym[idx_arr[2*iB+0]] == AS) ^ (B->sym[i] == AS)) ||
1558 1559 1560 1561
                idx_arr[2*idx_B[i+1]+0] == -1 ||
                idx_B[i+1] != idx_A[idx_arr[2*iB+0]+1]){
              sidx = 2*i+1;
              break;
1562 1563 1564 1565
            } else if (A->sym[idx_arr[2*iB+0]] == NS){
              sidx = 2*i;
              bsym = B->sym[i];
              break;
1566 1567 1568 1569 1570 1571 1572 1573
            }
          } else if (idx_arr[2*idx_B[i+1]+0] != -1){
            sidx = 2*i+1;
            break;
          }
        }
      }
    }
1574
    //if we have e.g. b[""] = A["ij"] with SY A, symmetry preserved bu t need to account for diagonal, this just unpacks (FIXME: suboptimal)
1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590
    if (sidx == -1){
      for (i=0; i<A->order; i++){
        if (A->sym[i] == SY){
          iA = idx_A[i];
          iA2 = idx_A[i+1];
          if (idx_arr[2*iA+1] == -1 &&
              idx_arr[2*iA2+1] == -1){
            sidx = 2*i;
            break;
          }
        }
      }
    } 
    if (nnew_sum != NULL && sidx != -1){
      if(sidx%2 == 0){
        new_sum->A = new tensor(A, 0, 0);
1591 1592
        int nA_sym[A->order];
        memcpy(nA_sym, new_sum->A->sym, sizeof(int)*new_sum->A->order);
1593
        nA_sym[sidx/2] = bsym;
1594
        new_sum->A->set_sym(nA_sym);
1595 1596
      } else {
        new_sum->B = new tensor(B, 0, 0);
1597 1598 1599 1600
        int nB_sym[B->order];
        memcpy(nB_sym, new_sum->B->sym, sizeof(int)*new_sum->B->order);
        nB_sym[sidx/2] = NS;
        new_sum->B->set_sym(nB_sym);
1601 1602
      }
    }
solomon's avatar
solomon committed
1603
    CTF_int::cdealloc(idx_arr);
1604 1605 1606 1607 1608 1609 1610 1611
    return sidx;
  }

  void summation::check_consistency(){
    int i, num_tot, len;
    int iA, iB;
    int * idx_arr;
       
1612 1613
    inv_idx(A->order, idx_A,
            B->order, idx_B,
1614 1615 1616 1617 1618 1619 1620 1621 1622 1623
            &num_tot, &idx_arr);

    for (i=0; i<num_tot; i++){
      len = -1;
      iA = idx_arr[2*i+0];
      iB = idx_arr[2*i+1];
      if (iA != -1){
        len = A->lens[iA];
      }
      if (len != -1 && iB != -1 && len != B->lens[iB]){
solomon's avatar
solomon committed
1624
        if (A->wrld->cdt.rank == 0){
1625
          printf("i = %d Error in sum call: The %dth edge length (%d) of tensor %s does not",
solomon's avatar
solomon committed
1626
                  i, iA, len, A->name);
1627
          printf("match the %dth edge length (%d) of tensor %s.\n",
solomon's avatar
solomon committed
1628
                  iB, B->lens[iB], B->name);
1629 1630 1631 1632
        }
        ABORT;
      }
    }
solomon's avatar
solomon committed
1633
    CTF_int::cdealloc(idx_arr);
1634 1635 1636 1637

  }


1638
  int summation::is_equal(summation const & os){
1639 1640
    int i;

1641 1642
    if (A != os.A) return 0;
    if (B != os.B) return 0;
1643

solomon's avatar
solomon committed
1644
    for (i=0; i<A->order; i++){
1645
      if (idx_A[i] != os.idx_A[i]) return 0;
1646
    }
solomon's avatar
solomon committed
1647
    for (i=0; i<B->order; i++){
1648
      if (idx_B[i] != os.idx_B[i]) return 0;
1649 1650 1651
    }
    return 1;
  }
1652

1653 1654
  int summation::check_mapping(){
    int i, pass, order_tot, iA, iB;
1655
    int * idx_arr; //, * phys_map;
1656
    //mapping * map;
1657 1658 1659 1660 1661 1662 1663 1664

    TAU_FSTART(check_sum_mapping);
    pass = 1;
    
    if (A->is_mapped == 0) pass = 0;
    if (B->is_mapped == 0) pass = 0;
    
    
1665 1666
    if (A->is_folded == 1) pass = 0;
    if (B->is_folded == 1) pass = 0;
1667 1668 1669 1670
    
    if (A->topo != B->topo) pass = 0;

    if (pass==0){
1671
      DPRINTF(4,"failed confirmation here\n");
1672 1673 1674 1675
      TAU_FSTOP(check_sum_mapping);
      return 0;
    }
    
1676 1677
    //CTF_int::alloc_ptr(sizeof(int)*A->topo->order, (void**)&phys_map);
    //memset(phys_map, 0, sizeof(int)*A->topo->order);
1678

1679 1680
    inv_idx(A->order, idx_A,
            B->order, idx_B,
1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698
            &order_tot, &idx_arr);

    if (!check_self_mapping(A, idx_A))
      pass = 0;
    if (!check_self_mapping(B, idx_B))
      pass = 0;
    if (pass == 0)
      DPRINTF(4,"failed confirmation here\n");

    for (i=0; i<order_tot; i++){
      iA = idx_arr[2*i];
      iB = idx_arr[2*i+1];
      if (iA != -1 && iB != -1) {
        if (!comp_dim_map(&A->edge_map[iA], &B->edge_map[iB])){
          pass = 0;
          DPRINTF(4,"failed confirmation here i=%d\n",i);
        }
      }
1699
      /*if (iA != -1) {
1700 1701 1702 1703 1704 1705 1706 1707 1708 1709 1710 1711 1712 1713 1714 1715 1716 1717
        map = &A->edge_map[iA];
        if (map->type == PHYSICAL_MAP)
          phys_map[map->cdt] = 1;
        while (map->has_child) {
          map = map->child;
          if (map->type == PHYSICAL_MAP)
            phys_map[map->cdt] = 1;
        }
      }
      if (iB != -1){
        map = &B->edge_map[iB];
        if (map->type == PHYSICAL_MAP)
          phys_map[map->cdt] = 1;
        while (map->has_child) {
          map = map->child;
          if (map->type == PHYSICAL_MAP)
            phys_map[map->cdt] = 1;
        }
1718
      }*/
1719 1720 1721 1722 1723 1724 1725 1726 1727 1728
    }
    /* Ensure that something is mapped to each dimension, since replciation
       does not make sense in sum for all tensors */
  /*  for (i=0; i<topovec[A->itopo].order; i++){
      if (phys_map[i] == 0) {
        pass = 0;
        DPRINTF(3,"failed confirmation here i=%d\n",i);
      }
    }*/

solomon's avatar
solomon committed
1729 1730
    //CTF_int::cdealloc(phys_map);
    CTF_int::cdealloc(idx_arr);
1731

1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746
    //if we have don't have an additive id we can't replicate
    if (B->sr->addid() == NULL){
      int ndim_mapped = 0;
      for (int j=0; j<B->order; j++){
        mapping * map = &B->edge_map[j];
        if (map->type == PHYSICAL_MAP) ndim_mapped++;
        while (map->has_child) {
          map = map->child;
          if (map->type == PHYSICAL_MAP)
            ndim_mapped++;
        }
        if (ndim_mapped < B->topo->order) pass = 0;
      }
    }
       
1747 1748 1749 1750 1751
    TAU_FSTOP(check_sum_mapping);

    return pass;
  }

1752 1753 1754
  int summation::map_sum_indices(topology const * topo){
    int tsr_order, isum, iA, iB, i, j, jsum, jX, stat;
    int * tsr_edge_len, * tsr_sym_table, * restricted;
solomon's avatar
solomon committed
1755
    int * idx_arr, * idx_sum;
1756 1757 1758 1759 1760
    int num_sum, num_tot, idx_num;
    idx_num = 2;
    mapping * sum_map;

    TAU_FSTART(map_sum_indices);
solomon's avatar
solomon committed
1761

1762 1763
    inv_idx(A->order, idx_A,
            B->order, idx_B,
solomon's avatar
solomon committed
1764 1765
            &num_tot, &idx_arr);

1766
    CTF_int::alloc_ptr(sizeof(int)*num_tot, (void**)&idx_sum);
solomon's avatar
solomon committed
1767 1768 1769 1770 1771 1772 1773
    
    num_sum = 0;
    for (i=0; i<num_tot; i++){
      if (idx_arr[2*i] != -1 && idx_arr[2*i+1] != -1){
        idx_sum[num_sum] = i;
        num_sum++;
      }
1774
    }
1775 1776 1777 1778

    tsr_order = num_sum;


1779 1780 1781 1782
    CTF_int::alloc_ptr(tsr_order*sizeof(int),           (void**)&restricted);
    CTF_int::alloc_ptr(tsr_order*sizeof(int),           (void**)&tsr_edge_len);
    CTF_int::alloc_ptr(tsr_order*tsr_order*sizeof(int), (void**)&tsr_sym_table);
    CTF_int::alloc_ptr(tsr_order*sizeof(mapping),       (void**)&sum_map);
1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838 1839 1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861

    memset(tsr_sym_table, 0, tsr_order*tsr_order*sizeof(int));
    memset(restricted, 0, tsr_order*sizeof(int));

    for (i=0; i<tsr_order; i++){ 
      sum_map[i].type             = NOT_MAPPED; 
      sum_map[i].has_child        = 0;
      sum_map[i].np               = 1;
    }
    for (i=0; i<num_sum; i++){
      isum = idx_sum[i];
      iA = idx_arr[isum*2+0];
      iB = idx_arr[isum*2+1];

      if (A->edge_map[iA].type != NOT_MAPPED){
        ASSERT(B->edge_map[iB].type == NOT_MAPPED);
        copy_mapping(1, &A->edge_map[iA], &sum_map[i]);
      } else if (B->edge_map[iB].type != NOT_MAPPED){
        copy_mapping(1, &B->edge_map[iB], &sum_map[i]);
      }
    }

    /* Map a tensor of dimension.
     * Set the edge lengths and symmetries according to those in sum dims of A and B.
     * This gives us a mapping for the common mapped dimensions of tensors A and B. */
    for (i=0; i<num_sum; i++){
      isum = idx_sum[i];
      iA = idx_arr[isum*idx_num+0];
      iB = idx_arr[isum*idx_num+1];

      tsr_edge_len[i] = A->pad_edge_len[iA];

      /* Check if A has symmetry among the dimensions being contracted over.
       * Ignore symmetry with non-contraction dimensions.
       * FIXME: this algorithm can be more efficient but should not be a bottleneck */
      if (A->sym[iA] != NS){
        for (j=0; j<num_sum; j++){
          jsum = idx_sum[j];
          jX = idx_arr[jsum*idx_num+0];
          if (jX == iA+1){
            tsr_sym_table[i*tsr_order+j] = 1;
            tsr_sym_table[j*tsr_order+i] = 1;
          }
        }
      }
      if (B->sym[iB] != NS){
        for (j=0; j<num_sum; j++){
          jsum = idx_sum[j];
          jX = idx_arr[jsum*idx_num+1];
          if (jX == iB+1){
            tsr_sym_table[i*tsr_order+j] = 1;
            tsr_sym_table[j*tsr_order+i] = 1;
          }
        }
      }
    }
    /* Run the mapping algorithm on this construct */
    stat = map_tensor(topo->order,        tsr_order, 
                      tsr_edge_len,       tsr_sym_table,
                      restricted,         topo->dim_comm,
                      NULL,               0,
                      sum_map);

    if (stat == ERROR){
      TAU_FSTOP(map_sum_indices);
      return ERROR;
    }
    
    /* define mapping of tensors A and B according to the mapping of sum dims */
    if (stat == SUCCESS){
      for (i=0; i<num_sum; i++){
        isum = idx_sum[i];
        iA = idx_arr[isum*idx_num+0];
        iB = idx_arr[isum*idx_num+1];

        copy_mapping(1, &sum_map[i], &A->edge_map[iA]);
        copy_mapping(1, &sum_map[i], &B->edge_map[iB]);
      }
    }
solomon's avatar
solomon committed
1862 1863 1864
    CTF_int::cdealloc(restricted);
    CTF_int::cdealloc(tsr_edge_len);
    CTF_int::cdealloc(tsr_sym_table);
1865 1866 1867
    for (i=0; i<num_sum; i++){
      sum_map[i].clear();
    }
solomon's avatar
solomon committed
1868 1869 1870
    CTF_int::cdealloc(sum_map);
    CTF_int::cdealloc(idx_sum);
    CTF_int::cdealloc(idx_arr);
1871 1872 1873 1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888

    TAU_FSTOP(map_sum_indices);
    return stat;

  }

  int summation::map(){
    int i, ret, need_remap;
    int need_remap_A, need_remap_B;
    int d;
    topology * old_topo_A, * old_topo_B;
    int btopo;
    int gtopo;

    ASSERT(A->wrld == B->wrld);
    World * wrld = A->wrld;
   
    TAU_FSTART(map_tensor_pair);
solomon's avatar
solomon committed
1889
  #if DEBUG >= 2
1890
    if (wrld->rank == 0)
solomon's avatar
solomon committed
1891
      printf("Initial mappings:\n");
1892
    A->print_map(stdout);
1893
    B->print_map(stdout);
solomon's avatar
solomon committed
1894 1895
  #endif

1896 1897 1898 1899 1900 1901
    //FIXME: try to avoid unfolding immediately, as its not always necessary
    A->unfold();
    B->unfold();
    A->set_padding();
    B->set_padding();

1902

1903 1904
    distribution dA(A);
    distribution dB(B);
1905 1906 1907 1908 1909 1910
    old_topo_A = A->topo;
    old_topo_B = B->topo;
    mapping * old_map_A = new mapping[A->order];
    mapping * old_map_B = new mapping[B->order];
    copy_mapping(A->order, A->edge_map, old_map_A);
    copy_mapping(B->order, B->edge_map, old_map_B);
solomon's avatar
solomon committed
1911
    btopo = -1;
1912 1913
    int64_t size;
    int64_t min_size = INT64_MAX;
solomon's avatar
solomon committed
1914
    /* Attempt to map to all possible permutations of processor topology */
1915
    for (i=A->wrld->cdt.rank; i<2*(int)A->wrld->topovec.size(); i+=A->wrld->cdt.np){
solomon's avatar
solomon committed
1916
  //  for (i=global_comm.rank*topovec.size(); i<2*(int)topovec.size(); i++){
1917 1918 1919 1920
      A->clear_mapping();
      B->clear_mapping();
      A->set_padding();
      B->set_padding();
solomon's avatar
solomon committed
1921

1922 1923 1924 1925
      A->topo = wrld->topovec[i/2];
      B->topo = wrld->topovec[i/2];
      A->is_mapped = 1;
      B->is_mapped = 1;
solomon's avatar
solomon committed
1926 1927

      if (i%2 == 0){
1928 1929
        ret = map_self_indices(A, idx_A);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1930 1931
        else if (ret != SUCCESS) return ret;
      } else {
1932 1933
        ret = map_self_indices(B, idx_B);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1934 1935
        else if (ret != SUCCESS) return ret;
      }
1936 1937
      ret = map_sum_indices(A->topo);
      if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1938 1939 1940 1941
      else if (ret != SUCCESS){
        return ret;
      }
      if (i%2 == 0){
1942 1943
        ret = map_self_indices(A, idx_A);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1944 1945
        else if (ret != SUCCESS) return ret;
      } else {
1946 1947
        ret = map_self_indices(B, idx_B);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1948 1949 1950 1951
        else if (ret != SUCCESS) return ret;
      }

      if (i%2 == 0){
1952 1953
        ret = map_self_indices(A, idx_A);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1954
        else if (ret != SUCCESS) return ret;
1955 1956 1957
        ret = A->map_tensor_rem(A->topo->order, 
                                A->topo->dim_comm);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1958
        else if (ret != SUCCESS) return ret;
1959 1960 1961 1962 1963 1964
        copy_mapping(A->order, B->order,
                     idx_A, A->edge_map, 
                     idx_B, B->edge_map,0);
        ret = B->map_tensor_rem(B->topo->order, 
                                B->topo->dim_comm);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1965 1966
        else if (ret != SUCCESS) return ret;
      } else {
1967 1968
        ret = map_self_indices(B, idx_B);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1969
        else if (ret != SUCCESS) return ret;
1970 1971 1972
        ret = B->map_tensor_rem(B->topo->order, 
                                B->topo->dim_comm);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1973
        else if (ret != SUCCESS) return ret;
1974 1975 1976 1977 1978 1979
        copy_mapping(B->order, A->order,
                     idx_B, B->edge_map, 
                     idx_A, A->edge_map,0);
        ret = A->map_tensor_rem(A->topo->order, 
                                A->topo->dim_comm);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1980 1981 1982
        else if (ret != SUCCESS) return ret;
      }
      if (i%2 == 0){
1983 1984
        ret = map_self_indices(B, idx_B);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1985 1986
        else if (ret != SUCCESS) return ret;
      } else {
1987 1988
        ret = map_self_indices(A, idx_A);
        if (ret == NEGATIVE) continue;
solomon's avatar
solomon committed
1989 1990 1991
        else if (ret != SUCCESS) return ret;
      }

1992 1993
  /*    ret = map_symtsr(A->order, A->sym_table, A->edge_map);
      ret = map_symtsr(B->order, B->sym_table, B->edge_map);
solomon's avatar
solomon committed
1994 1995 1996
      if (ret!=SUCCESS) return ret;
      return SUCCESS;*/

1997
  #if DEBUG >= 4
1998 1999
      A->print_map(stdout,0);
      B->print_map(stdout,0);
solomon's avatar
solomon committed
2000
  #endif
2001 2002 2003 2004
      if (!check_mapping()) continue;
      A->set_padding();
      B->set_padding();
      size = A->size + B->size;
solomon's avatar
solomon committed
2005 2006 2007 2008

      need_remap_A = 0;
      need_remap_B = 0;

2009 2010 2011
      if (A->topo == old_topo_A){
        for (d=0; d<A->order; d++){
          if (!comp_dim_map(&A->edge_map[d],&old_map_A[d]))
solomon's avatar
solomon committed
2012 2013 2014 2015 2016
            need_remap_A = 1;
        }
      } else
        need_remap_A = 1;
      if (need_remap_A){
2017 2018
        if (can_block_reshuffle(A->order, dA.phase, A->edge_map)){
          size += A->size*log2(wrld->cdt.np);
solomon's avatar
solomon committed
2019
        } else {
2020
          size += 5.*A->size*log2(wrld->cdt.np);
solomon's avatar
solomon committed
2021 2022
        }
      }
2023 2024 2025
      if (B->topo == old_topo_B){
        for (d=0; d<B->order; d++){
          if (!comp_dim_map(&B->edge_map[d],&old_map_B[d]))
solomon's avatar
solomon committed
2026 2027 2028 2029 2030
            need_remap_B = 1;
        }
      } else
        need_remap_B = 1;
      if (need_remap_B){
2031 2032
        if (can_block_reshuffle(B->order, dB.phase, B->edge_map)){
          size += B->size*log2(wrld->cdt.np);
solomon's avatar
solomon committed
2033
        } else {
2034
          size += 5.*B->size*log2(wrld->cdt.np);
solomon's avatar
solomon committed
2035 2036 2037
        }
      }

2038 2039
      /*nvirt = (int64_t)calc_nvirt(A);
      tnvirt = nvirt*(int64_t)calc_nvirt(B);
solomon's avatar
solomon committed
2040 2041 2042 2043 2044 2045 2046 2047 2048 2049 2050 2051
      if (tnvirt < nvirt) nvirt = UINT64_MAX;
      else nvirt = tnvirt;
      if (btopo == -1 || nvirt < bnvirt ) {
        bnvirt = nvirt;
        btopo = i;      
      }*/
      if (btopo == -1 || size < min_size){
        min_size = size;
        btopo = i;      
      }
    }
    if (btopo == -1)
2052
      min_size = INT64_MAX;
solomon's avatar
solomon committed
2053
    /* pick lower dimensional mappings, if equivalent */
2054
    gtopo = get_best_topo(min_size, btopo, wrld->cdt);
solomon's avatar
solomon committed
2055 2056 2057 2058 2059 2060 2061
    TAU_FSTOP(map_tensor_pair);
    if (gtopo == -1){
      printf("ERROR: Failed to map pair!\n");
      ABORT;
      return ERROR;
    }
    
2062 2063 2064 2065
    A->clear_mapping();
    B->clear_mapping();
    A->set_padding();
    B->set_padding();
2066

2067 2068
    A->topo = wrld->topovec[gtopo/2];
    B->topo = wrld->topovec[gtopo/2];
2069 2070 2071

    A->is_mapped = 1;
    B->is_mapped = 1;
solomon's avatar
solomon committed
2072 2073
      
    if (gtopo%2 == 0){
2074
      ret = map_self_indices(A, idx_A);
solomon's avatar
solomon committed
2075
      ASSERT(ret == SUCCESS);
2076
    } else {
2077
      ret = map_self_indices(B, idx_B);
solomon's avatar
solomon committed
2078
      ASSERT(ret == SUCCESS);
2079
    }
2080
    ret = map_sum_indices(A->topo);
solomon's avatar
solomon committed
2081
    ASSERT(ret == SUCCESS);
2082 2083 2084 2085 2086 2087 2088
    if (gtopo%2 == 0){
      ret = map_self_indices(A, idx_A);
      ASSERT(ret == SUCCESS);
    } else {
      ret = map_self_indices(B, idx_B);
      ASSERT(ret == SUCCESS);
    }
2089

solomon's avatar
solomon committed
2090
    if (gtopo%2 == 0){
2091
      ret = map_self_indices(A, idx_A);
solomon's avatar
solomon committed
2092
      ASSERT(ret == SUCCESS);
2093 2094
      ret = A->map_tensor_rem(A->topo->order, 
                              A->topo->dim_comm);
solomon's avatar
solomon committed
2095
      ASSERT(ret == SUCCESS);
2096 2097 2098 2099 2100
      copy_mapping(A->order, B->order,
                   idx_A, A->edge_map, 
                   idx_B, B->edge_map,0);
      ret = B->map_tensor_rem(B->topo->order, 
                              B->topo->dim_comm);
solomon's avatar
solomon committed
2101
      ASSERT(ret == SUCCESS);
2102
    } else {
2103
      ret = map_self_indices(B, idx_B);
solomon's avatar
solomon committed
2104
      ASSERT(ret == SUCCESS);
2105 2106
      ret = B->map_tensor_rem(B->topo->order, 
                              B->topo->dim_comm);
solomon's avatar
solomon committed
2107
      ASSERT(ret == SUCCESS);
2108 2109 2110 2111 2112
      copy_mapping(B->order, A->order,
                   idx_B, B->edge_map, 
                   idx_A, A->edge_map,0);
      ret = A->map_tensor_rem(A->topo->order, 
                              A->topo->dim_comm);
solomon's avatar
solomon committed
2113
      ASSERT(ret == SUCCESS);
2114
    }
2115
    if (gtopo%2 == 0){
2116
      ret = map_self_indices(B, idx_B);
2117 2118
      ASSERT(ret == SUCCESS);
    } else {
2119
      ret = map_self_indices(A, idx_A);
2120 2121 2122
      ASSERT(ret == SUCCESS);
    }

2123
    ASSERT(check_mapping());
2124

2125 2126
    A->set_padding();
    B->set_padding();
2127
  #if DEBUG > 2
2128
    if (wrld->cdt.rank == 0)
solomon's avatar
solomon committed
2129
      printf("New mappings:\n");
2130 2131
    A->print_map(stdout);
    B->print_map(stdout);
solomon's avatar
solomon committed
2132
  #endif
2133

solomon's avatar
solomon committed
2134 2135
    TAU_FSTART(redistribute_for_sum);
   
2136 2137
    A->is_cyclic = 1;
    B->is_cyclic = 1;
solomon's avatar
solomon committed
2138
    need_remap = 0;
2139 2140 2141
    if (A->topo == old_topo_A){
      for (d=0; d<A->order; d++){
        if (!comp_dim_map(&A->edge_map[d],&old_map_A[d]))
solomon's avatar
solomon committed
2142
          need_remap = 1;
2143 2144
      }
    } else
solomon's avatar
solomon committed
2145 2146
      need_remap = 1;
    if (need_remap)
2147
      A->redistribute(dA);
solomon's avatar
solomon committed
2148
    need_remap = 0;
2149 2150 2151
    if (B->topo == old_topo_B){
      for (d=0; d<B->order; d++){
        if (!comp_dim_map(&B->edge_map[d],&old_map_B[d]))
solomon's avatar
solomon committed
2152
          need_remap = 1;
2153 2154
      }
    } else
solomon's avatar
solomon committed
2155 2156
      need_remap = 1;
    if (need_remap)
2157
      B->redistribute(dB);
solomon's avatar
solomon committed
2158 2159

    TAU_FSTOP(redistribute_for_sum);
2160 2161
    delete [] old_map_A;
    delete [] old_map_B;
2162

solomon's avatar
solomon committed
2163
    return SUCCESS;
2164
  }
2165 2166 2167 2168 2169 2170 2171 2172

  void summation::print(){
    int i,j,max,ex_A, ex_B;
    max = A->order+B->order;
    CommData global_comm = A->wrld->cdt;
    MPI_Barrier(global_comm.cm);
    if (global_comm.rank == 0){
      printf("Summing Tensor %s into %s\n", A->name, B->name);
2173 2174 2175 2176 2177 2178 2179
      printf("alpha is "); 
      if (alpha != NULL) A->sr->print(alpha);
      else printf("NULL");
      printf("\nbeta is "); 
      if (beta != NULL) B->sr->print(beta);
      else printf("NULL"); 
      printf("\n");
2180 2181 2182 2183 2184 2185 2186 2187 2188
      printf("Summation index table:\n");
      printf("     A     B\n");
      for (i=0; i<max; i++){
        ex_A=0;
        ex_B=0;
        printf("%d:   ",i);
        for (j=0; j<A->order; j++){
          if (idx_A[j] == i){
            ex_A++;
solomon's avatar
solomon committed
2189 2190 2191 2192 2193 2194
            if (A->sym[j] == SY)
              printf("%dY ",j);
            else if (A->sym[j] == SH)
              printf("%dH ",j);
            else if (A->sym[j] == AS)
              printf("%dS ",j);
2195 2196 2197 2198 2199 2200 2201 2202 2203 2204 2205
            else
              printf("%d  ",j);
          }
        }
        if (ex_A == 0)
          printf("      ");
        if (ex_A == 1)
          printf("   ");
        for (j=0; j<B->order; j++){
          if (idx_B[j] == i){
            ex_B=1;
solomon's avatar
solomon committed
2206 2207 2208 2209 2210 2211
            if (B->sym[j] == SY)
              printf("%dY ",j);
            else if (B->sym[j] == SH)
              printf("%dH ",j);
            else if (B->sym[j] == AS)
              printf("%dS ",j);
2212 2213 2214 2215 2216 2217 2218 2219 2220
            else
              printf("%d  ",j);
          }
        }
        printf("\n");
        if (ex_A + ex_B== 0) break;
      }
    }
  }
2221
              
2222

2223 2224 2225 2226 2227 2228 2229 2230 2231 2232 2233 2234 2235 2236 2237
  void summation::sp_sum(){
    int64_t num_pair;
    char * mapped_data;
    
    bool is_idx_matched = true;
    if (A->order != B->order)
      is_idx_matched = false;
    else {
      for (int o=0; o<A->order; o++){
        if (idx_A[o] != idx_B[o]){
          is_idx_matched = false;
        }
      }
    }

2238

2239
    //read data from A    
2240
    A->read_local_nnz(&num_pair, &mapped_data);
2241 2242 2243 2244 2245 2246 2247 2248 2249 2250 2251 2252 2253

    if (!is_idx_matched){
      int64_t lda_A[A->order];
      int64_t lda_B[B->order];
      lda_A[0] = 1;
      for (int o=1; o<A->order; o++){
        lda_A[o] = lda_A[o-1]*A->lens[o];
      }
      lda_B[0] = 1;
      for (int o=1; o<B->order; o++){
        lda_B[o] = lda_B[o-1]*B->lens[o];
      }
      PairIterator pi(A->sr, mapped_data);
2254
#ifdef USE_OMP
2255
      #pragma omp parallel for
2256
#endif
2257 2258 2259 2260 2261 2262 2263 2264 2265 2266 2267 2268 2269 2270 2271
      for (int i=0; i<num_pair; i++){
        int64_t k = pi[i].k();
        int64_t k_new = 0;
        for (int o=0; o<A->order; o++){
          int64_t kpart = (k/lda_A[o])%A->lens[o];
          //FIXME: slow, but handles diagonal indexing, probably worth having separate versions
          for (int q=0; q<B->order; q++){
            if (idx_A[o] == idx_B[q]){
              k_new += kpart*lda_B[q];
            }
          }
        }
        ((int64_t*)(pi[i].ptr))[0] = k_new;
      }

2272 2273 2274 2275 2276 2277 2278 2279 2280 2281 2282 2283 2284 2285 2286 2287 2288 2289 2290 2291 2292 2293 2294 2295 2296 2297 2298 2299 2300 2301 2302 2303 2304 2305 2306 2307 2308 2309 2310 2311 2312 2313 2314 2315 2316 2317 2318 2319 2320 2321 2322 2323 2324 2325 2326 2327
      // when idx_A has indices idx_B does not, we need to reduce, which can be done partially here since the elements of A should be sorted
      bool is_reduce = false;
      for (int oA=0; oA<A->order; oA++){
        bool inB = false;
        for (int oB=0; oB<B->order; oB++){
          if (idx_A[oA] == idx_B[oB]){
            inB = true;
          }
        }
        if (!inB) is_reduce = true;
      }
  
      if (is_reduce && num_pair > 0){
        pi.sort(num_pair);
        int64_t nuniq=1;
        for (int64_t i=1; i<num_pair; i++){
          if (pi[i].k() != pi[i-1].k()) nuniq++;
        }
        if (nuniq != num_pair){
          char * swap_data = mapped_data;
          alloc_ptr(A->sr->pair_size()*nuniq, (void**)&mapped_data);
          PairIterator pi_new(A->sr, mapped_data);
          int64_t cp_st = 0;
          int64_t acc_st = -1;
          int64_t pfx = 0;
          for (int64_t i=1; i<num_pair; i++){
            if (pi[i].k() == pi[i-1].k()){
              if (cp_st < i){ 
                memcpy(pi_new[pfx].ptr, pi[cp_st].ptr, A->sr->pair_size()*(i-cp_st));
                pfx += i-cp_st;
              }
              cp_st = i+1;

              if (acc_st == -1) acc_st = i;
            } else {
              if (acc_st != -1){
                for (int64_t j=acc_st; j<i; j++){
                  A->sr->add(pi_new[pfx-1].d(), pi[j].d(), pi_new[pfx-1].d());
                }
              }
              acc_st = -1;
            }           
          }
          if (cp_st < num_pair)
            memcpy(pi_new[pfx].ptr, pi[cp_st].ptr, A->sr->pair_size()*(num_pair-cp_st));
          if (acc_st != -1){
            for (int64_t j=acc_st; j<num_pair; j++){
              A->sr->add(pi_new[pfx-1].d(), pi[j].d(), pi_new[pfx-1].d());
            }
          }
          cdealloc(swap_data);
          num_pair = nuniq;
        }
      }

      // if applying custom function, apply immediately on reduced form
2328
      if (is_custom && !func->is_accumulator()){
2329 2330 2331 2332 2333 2334 2335 2336
        char * swap_data = mapped_data;
        alloc_ptr(B->sr->pair_size()*num_pair, (void**)&mapped_data);
        PairIterator pi_new(B->sr, mapped_data);
#ifdef USE_OMP
        #pragma omp parallel for
#endif
        for (int64_t i=0; i<num_pair; i++){
          if (alpha == NULL)
2337
            func->apply_f(pi[i].d(), pi_new[i].d());
2338 2339 2340
          else  {
            char tmp_A[A->sr->el_size];
            A->sr->mul(pi[i].d(), alpha, tmp_A);
2341
            func->apply_f(tmp_A, pi_new[i].d());
2342 2343 2344 2345 2346 2347 2348 2349 2350 2351 2352 2353 2354 2355 2356 2357 2358 2359 2360 2361 2362 2363 2364 2365 2366 2367 2368 2369 2370 2371 2372 2373 2374 2375 2376 2377 2378 2379 2380 2381 2382 2383 2384 2385 2386 2387 2388 2389 2390 2391 2392 2393 2394 2395 2396 2397 2398 2399 2400 2401 2402 2403 2404 2405 2406 2407 2408 2409 2410 2411 2412 2413 2414
          }
        }
        cdealloc(swap_data);
        alpha = NULL;
      }
  
      // when idx_B has indices idx_A does not, we need to map, which we do by replicating the key value pairs of B
      // FIXME this is probably not most efficient, but not entirely stupid, as at least the set of replicated pairs is not expected to be bigger than B
      int nmap_idx = 0;
      int64_t map_idx_len[B->order];
      int64_t map_idx_lda[B->order];
      int map_idx_rev[B->order];
      for (int oB=0; oB<B->order; oB++){
        bool inA = false;
        for (int oA=0; oA<A->order; oA++){
          if (idx_A[oA] == idx_B[oB]){
            inA = true;
          }
        }
        if (!inA){ 
          bool is_rep=false;
          for (int ooB=0; ooB<oB; ooB++){
            if (idx_B[ooB] == idx_B[oB]){
              is_rep = true;
              map_idx_lda[map_idx_rev[ooB]] += lda_B[oB];
              break;
            }
          }
          if (!is_rep){
            map_idx_len[nmap_idx] = B->lens[oB];
            map_idx_lda[nmap_idx] = lda_B[oB];
            map_idx_rev[nmap_idx] = oB;
            nmap_idx++;
          }
        }
      }
      if (nmap_idx > 0){
        int64_t tot_rep=1;
        for (int midx=0; midx<nmap_idx; midx++){
          tot_rep *= map_idx_len[midx];
        }
        char * swap_data = mapped_data;
        alloc_ptr(A->sr->pair_size()*num_pair*tot_rep, (void**)&mapped_data);
        PairIterator pi_new(A->sr, mapped_data);
#ifdef USE_OMP
        #pragma omp parallel for
#endif
        for (int64_t i=0; i<num_pair; i++){
          for (int64_t r=0; r<tot_rep; r++){
            memcpy(pi_new[i*tot_rep+r].ptr, pi[i].ptr, A->sr->pair_size());
          }
        }
#ifdef USE_OMP
        #pragma omp parallel for
#endif
        for (int64_t i=0; i<num_pair; i++){
          int64_t phase=1;
          for (int midx=0; midx<nmap_idx; midx++){
            int64_t stride=phase;
            phase *= map_idx_len[midx];
            for (int64_t r=0; r<tot_rep/phase; r++){
              for (int64_t m=1; m<map_idx_len[midx]; m++){
                for (int64_t s=0; s<stride; s++){
                  ((int64_t*)(pi_new[i*tot_rep + r*phase + m*stride + s].ptr))[0] += m*map_idx_lda[midx];
                }
              }
            }
          }
        }
        cdealloc(swap_data);
        num_pair *= tot_rep;
      }
    }
2415 2416 2417 2418 2419 2420
    
    B->write(num_pair, alpha, beta, mapped_data, 'w');
    cdealloc(mapped_data);

  }

2421
}