tensorflow/core/kernels/roll_op.cc
// Use memcpy to copy memory in groups when the data type supports memcpy
Use memcpy to copy memory in groups when the data type supports memcpy
Use memcpy to copy memory in groups when the data type supports memcpy
void DoRollWithMemcpy(const OpKernelContext* context,
                      const int64_t num_elements, const int num_dims,
                      const gtl::ArraySlice<int32> dim_size, const T* input,
                      T* output, const gtl::ArraySlice<int32> threshold,
                      const gtl::ArraySlice<int64_t> dim_range,
                      const int64_t isd) {
  auto work = [input, output, num_dims, &dim_size, &threshold, &dim_range, isd](
                  int64_t start, int64_t end) {
    const int64_t isd_range = std::max<int>(dim_range[isd], 1);
    const int64_t isd_stride = isd_range / std::max<int>(dim_size[isd], 1);
    const int64_t start_remainder = (start % 2) * threshold[isd] * isd_stride;
    const int64_t end_remainder = (end % 2) * threshold[isd] * isd_stride;
    start = (start / 2) * isd_range + start_remainder;
    end = (end / 2) * isd_range + end_remainder;
    const T* in_ptr = &input[0];
    T* out_ptr = &output[0];
    in_ptr += start;
    out_ptr += start;
    gtl::InlinedVector<int, 4> indices(num_dims);
    int64_t remainder_offset = 0;
    for (int i = 0; i < num_dims; i++) {
      const int64_t stride = dim_range[i] / dim_size[i];
      const int shift = dim_size[i] - threshold[i];
      const int indx = (start / stride) % dim_size[i];
      indices[i] = indx;
      int out_indx = (indx + shift) % dim_size[i];
      if (i > isd) {
        out_indx = 0;
        remainder_offset += (out_indx - indx) * stride;
      }
      out_ptr += (out_indx - indx) * stride;
    }
    for (int i = num_dims - 1; i > isd; i--) indices[i] = 0;
    int isd_indx_skip = 0;
    int64_t group_size = 0;
    if (indices[isd] < threshold[isd]) {
      isd_indx_skip = threshold[isd] - indices[isd];
      group_size = isd_indx_skip * isd_stride + remainder_offset;
    } else {
      isd_indx_skip = dim_size[isd] - indices[isd];
      group_size = isd_indx_skip * isd_stride + remainder_offset;
    }
    int64_t i = start;
    while (i < end) {
      memcpy(out_ptr, in_ptr, group_size * sizeof(T));
      i += group_size;
      out_ptr += group_size;
      in_ptr += group_size;
      for (int j = isd; j >= 0; j--) {
        int inc = 1;
        if (j == isd) inc = isd_indx_skip;
        const int indx = (indices[j] + inc) % dim_size[j];
        indices[j] = indx;
        if (indx != 0) {
          if (indx == threshold[j]) {
            out_ptr -= dim_range[j];  
          }
          break;                         
        } else if (threshold[j] != 0) {  
          out_ptr += dim_range[j];       
        }
      }
      if (indices[isd] < threshold[isd]) {
        isd_indx_skip = threshold[isd] - indices[isd];
        group_size = isd_indx_skip * isd_stride;
      } else {
        isd_indx_skip = dim_size[isd] - indices[isd];
        group_size = isd_indx_skip * isd_stride;
      }
    }
  };
  auto worker_threads = context->device()->tensorflow_cpu_worker_threads();
  const int64_t ave_group_size = dim_range[isd] / 2;
  const int total_work = 2 * num_elements / std::max<int>(dim_range[isd], 1);
  const int cost_per_group = 25000 * sizeof(T) * ave_group_size;
  Shard(worker_threads->num_threads, worker_threads->workers, total_work,
        cost_per_group, std::move(work));
}
void DoRollWithMemcpy ( const OpKernelContext * context , const int64_t num_elements , const int num_dims , const gtl :: ArraySlice < int32 > dim_size , const T * input , T * output , const gtl :: ArraySlice < int32 > threshold , const gtl :: ArraySlice < int64_t > dim_range , const int64_t isd ) { auto work = [ input , output , num_dims , & dim_size , & threshold , & dim_range , isd ] ( int64_t start , int64_t end ) { const int64_t isd_range = std :: max < int > ( dim_range [ isd ] , 1 ) ; const int64_t isd_stride = isd_range / std :: max < int > ( dim_size [ isd ] , 1 ) ; const int64_t start_remainder = ( start % 2 ) * threshold [ isd ] * isd_stride ; const int64_t end_remainder = ( end % 2 ) * threshold [ isd ] * isd_stride ; start = ( start / 2 ) * isd_range + start_remainder ; end = ( end / 2 ) * isd_range + end_remainder ; const T * in_ptr = & input [ 0 ] ; T * out_ptr = & output [ 0 ] ; in_ptr += start ; out_ptr += start ; gtl :: InlinedVector < int , 4 > indices ( num_dims ) ; int64_t remainder_offset = 0 ; for ( int i = 0 ; i < num_dims ; i ++ ) { const int64_t stride = dim_range [ i ] / dim_size [ i ] ; const int shift = dim_size [ i ] - threshold [ i ] ; const int indx = ( start / stride ) % dim_size [ i ] ; indices [ i ] = indx ; int out_indx = ( indx + shift ) % dim_size [ i ] ; if ( i > isd ) { out_indx = 0 ; remainder_offset += ( out_indx - indx ) * stride ; } out_ptr += ( out_indx - indx ) * stride ; } for ( int i = num_dims - 1 ; i > isd ; i -- ) indices [ i ] = 0 ; int isd_indx_skip = 0 ; int64_t group_size = 0 ; if ( indices [ isd ] < threshold [ isd ] ) { isd_indx_skip = threshold [ isd ] - indices [ isd ] ; group_size = isd_indx_skip * isd_stride + remainder_offset ; } else { isd_indx_skip = dim_size [ isd ] - indices [ isd ] ; group_size = isd_indx_skip * isd_stride + remainder_offset ; } int64_t i = start ; while ( i < end ) { memcpy ( out_ptr , in_ptr , group_size * sizeof ( T ) ) ; i += group_size ; out_ptr += group_size ; in_ptr += group_size ; for ( int j = isd ; j >= 0 ; j -- ) { int inc = 1 ; if ( j == isd ) inc = isd_indx_skip ; const int indx = ( indices [ j ] + inc ) % dim_size [ j ] ; indices [ j ] = indx ; if ( indx != 0 ) { if ( indx == threshold [ j ] ) { out_ptr -= dim_range [ j ] ; } break ; } else if ( threshold [ j ] != 0 ) { out_ptr += dim_range [ j ] ; } } if ( indices [ isd ] < threshold [ isd ] ) { isd_indx_skip = threshold [ isd ] - indices [ isd ] ; group_size = isd_indx_skip * isd_stride ; } else { isd_indx_skip = dim_size [ isd ] - indices [ isd ] ; group_size = isd_indx_skip * isd_stride ; } } } ; auto worker_threads = context -> device ( ) -> tensorflow_cpu_worker_threads ( ) ; const int64_t ave_group_size = dim_range [ isd ] / 2 ; const int total_work = 2 * num_elements / std :: max < int > ( dim_range [ isd ] , 1 ) ; const int cost_per_group = 25000 * sizeof ( T ) * ave_group_size ; Shard ( worker_threads -> num_threads , worker_threads -> workers , total_work , cost_per_group , std :: move ( work ) ) ; }
Use memcpy to copy memory in groups when the data type supports memcpy
Use memcpy to copy memory in groups when the data type supports memcpy
// the number of indices over in the flattened tensor you need to skip in // order to make it over from one side of the isd to the other // the distance along the flattend tensor to the next element in the isd // start and end represent the i-th group currently so we will convert // them into numbers representing the i-th elements. // there are 2 groups per isd one for all elements before threshold[isd] // and another for all elements after threshold[isd]. // array of indices for each dimension // indices = [i, j, k, l, m, n] // the offset needed to make all inner non-shifting dimensions become 0 // initialize indices // stride is the number of indices over in the flattened tensor // you need to skip in order to make it over to an adjacent element // along a dimension. dim_size[i] != 0 because we set it to max(dim, 1) // calculate dimension index after the shift // trailing zeroes for indices after the inner shifted dimension // set trailing zeroes for indices after the inner shifted dimension // the number of indices in the isd dimension the next group will skip // to make it to the next threshold or end point // the size of the next group // initialize isd_indx_skip and group_size // copy group of elements // shift i and the pointers over to the next group position // produce next combination of indices and adjust the out_ptr position // to fix the offset if necessary // the isd (inner shift dim) should skip to next threshold or endpoint // all dimensions to the left increment by 1 when a digit is carried // all dimensions to the right remain set to 0 //            +1 +1 +1 +isd_indx_skip // indices = [i, j, k, l, 0, 0] //                      ^isd // now wraps around // indx != 0 don't need to carry // if threshold is 0 shift is 0 // indx became 0 so reverse wrap around // set isd_indx_skip and group_size for next iteration // Shard // 25000 - experimentally determined with float and bool types