text: /******************************************************************************
  * Copyright (c) Intel Corporation - All rights reserved.                      *
  * This file is part of the LIBXSMM library.                                   *
  *                                                                             *
  * For information on the license, see the LICENSE file.                       *
  * Further information: https://github.com/hfp/libxsmm/                        *
  * SPDX-License-Identifier: BSD-3-Clause                                       *
  ******************************************************************************/
  /* Alexander Heinecke, Sasikanth Avancha (Intel Corp.)
  ******************************************************************************/
  
  #if defined(LIBXSMM_DNN_POOLING_FWD_BF16)
  # define _mm512_load_act(A)     _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepi16_epi32(_mm256_loadu_si256((__m256i*)(A))),16))
  #if 1
  # define _mm512_roundbf16rne(A) LIBXSMM_INTRINSICS_MM512_ROUNDNE_BF16(A)
  # define _mm512_stream_act(A,B) _mm256_stream_si256((__m256i*)(A),_mm512_cvtepi32_epi16(_mm512_srai_epi32(_mm512_roundbf16rne((B)),16)))
  # define _mm512_store_act(A,B)  _mm256_storeu_si256((__m256i*)(A),_mm512_cvtepi32_epi16(_mm512_srai_epi32(_mm512_roundbf16rne((B)),16)))
  #else
  # define _mm512_stream_act(A,B) _mm256_stream_si256((__m256i*)(A),_mm512_cvtepi32_epi16(_mm512_srai_epi32(_mm512_castps_si512((B)),16)))
  # define _mm512_store_act(A,B)  _mm256_storeu_si256((__m256i*)(A),_mm512_cvtepi32_epi16(_mm512_srai_epi32(_mm512_castps_si512((B)),16)))
  #endif
  #else
  # define _mm512_load_act(A)     _mm512_loadu_ps(A)
  # define _mm512_stream_act(A,B) LIBXSMM_INTRINSICS_MM512_STREAM_PS(A,B)
  # define _mm512_store_act(A,B)  _mm512_storeu_ps(A,B)
  #endif
  
  /* size variables, all const */
  const int nImg = handle->desc.N;
  const int ifh = handle->desc.H;
  const int ifw = handle->desc.W;
  const int sh = handle->desc.u;
  const int sw = handle->desc.v;
  const int ofh = handle->ofh;
  const int ofw = handle->ofw;
  const int iph = handle->desc.pad_h_in;
  const int ipw = handle->desc.pad_w_in;
  const int oph = handle->desc.pad_h_out;
  const int opw = handle->desc.pad_w_out;
  const int ofhp = ofh + 2*oph;
  const int ofwp = ofw + 2*opw;
  const int ifhp = ifh + 2*iph;
  const int ifwp = ifw + 2*ipw;
  /* here we assume that input and output blocking is similar */
  const int nBlocksFm = handle->blocksifm;
  
  /* computing first logical thread */
  const int ltid = tid - start_thread;
  /* number of tasks that could be run in parallel */
  const int work = nImg * nBlocksFm;
  /* compute chunk size */
  const int chunksize = (work % handle->desc.threads == 0) ? (work / handle->desc.threads) : ((work / handle->desc.threads) + 1);
  /* compute thr_begin and thr_end */
  const int thr_begin = (ltid * chunksize < work) ? (ltid * chunksize) : work;
  const int thr_end = ((ltid + 1) * chunksize < work) ? ((ltid + 1) * chunksize) : work;
  
  /* loop variables */
  int img = 0;
  int fm = 0;
  int imgfm = 0;
  int ho = 0;
  int wo = 0;
  int hi = 0;
  int wi = 0;
  int kh = 0;
  int kw = 0;
  int v = 0;
  #if defined(LIBXSMM_DNN_POOLING_FWD_AVG)
  #if defined(LIBXSMM_DNN_POOLING_FWD_BF16)
  float recp_pool_size = 1.0f/((float)handle->desc.R*(float)handle->desc.S);
  #else
  element_output_type recp_pool_size = 1.0f/((element_output_type)handle->desc.R*(element_output_type)handle->desc.S);
  #endif
  #endif
  
  /* multi-dim arrays declaration */
  #if defined(LIBXSMM_DNN_POOLING_FWD_BF16)
  float* lcl_buffer_ptr = ((float*)handle->scratch)+((size_t)ofh*(size_t)ofw*(size_t)64*(size_t)ltid);
  LIBXSMM_VLA_DECL(3,       float, lcl_output, lcl_buffer_ptr,                                                   ofw, 64);
  #else
  element_output_type* lcl_buffer_ptr = ((element_output_type*)handle->scratch)+((size_t)ofh*(size_t)ofw*(size_t)64*(size_t)ltid);
  LIBXSMM_VLA_DECL(3,       element_output_type, lcl_output, lcl_buffer_ptr,                                                   ofw, 64);
  #endif
  LIBXSMM_VLA_DECL(5, const element_input_type,  input,      (element_input_type* )handle->reg_input->data,  nBlocksFm, ifhp, ifwp, 64);
  LIBXSMM_VLA_DECL(5,       element_output_type, output,     (element_output_type*)handle->reg_output->data, nBlocksFm, ofhp, ofwp, 64);
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
  LIBXSMM_VLA_DECL(5,       element_mask_type,   mask,       (element_mask_type*  )handle->mask->data,       nBlocksFm,  ofh,  ofw, 64);
  #endif
  
  /* lazy barrier init */
  libxsmm_barrier_init(handle->barrier, ltid);
  
  for (imgfm = thr_begin; imgfm < thr_end; ++imgfm) {
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
    __m512i lcl_viadd  = _mm512_set_epi32( 15, 14, 13, 12, 11, 10,  9,  8,  7,  6,  5,  4,  3,  2,  1,  0 );
  #endif
    img = imgfm / nBlocksFm;
    fm = imgfm % nBlocksFm;
  
    for( v = 0; v < ofh*ofw*64; v+=16 ) {
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
      _mm512_storeu_ps( &(lcl_buffer_ptr[v]), _mm512_set1_ps(-FLT_MAX) );
  #endif
  #if defined(LIBXSMM_DNN_POOLING_FWD_AVG)
      _mm512_storeu_ps( &(lcl_buffer_ptr[v]), _mm512_setzero_ps() );
  #endif
    }
  
    for( ho = oph; ho < (ofh+oph); ho++ ) {
      hi = ((ho-oph) * sh) - handle->desc.pad_h;
      for( wo = opw; wo < (ofw+opw); wo++ ) {
        float*               lcl_output_ptr = &LIBXSMM_VLA_ACCESS(3, lcl_output, ho-oph, wo-opw, 0, ofw, 64);
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
        __m512i lcl_vmask  = _mm512_loadu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw,  0, nBlocksFm, ofh, ofw, 64) );
        __m512i lcl_vmask2 = _mm512_loadu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 16, nBlocksFm, ofh, ofw, 64) );
        __m512i lcl_vmask3 = _mm512_loadu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 32, nBlocksFm, ofh, ofw, 64) );
        __m512i lcl_vmask4 = _mm512_loadu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 48, nBlocksFm, ofh, ofw, 64) );
  #endif
        __m512 lcl_voutput  = _mm512_loadu_ps( lcl_output_ptr );
        __m512 lcl_voutput2 = _mm512_loadu_ps( lcl_output_ptr+16 );
        __m512 lcl_voutput3 = _mm512_loadu_ps( lcl_output_ptr+32 );
        __m512 lcl_voutput4 = _mm512_loadu_ps( lcl_output_ptr+48 );
  
        wi = ((wo-opw) * sw) - handle->desc.pad_w;
        for( kh = 0; kh < handle->desc.R; kh++ ) {
          if (hi+kh < 0 || hi+kh >= ifh) continue;
          for( kw = 0; kw < handle->desc.S; kw++ ) {
            if (wi+kw < 0 || wi+kw >= ifw) {
              continue;
            } else {
              const element_input_type*      input_ptr  = &LIBXSMM_VLA_ACCESS(5, input,      img, fm, hi+kh+iph, wi+kw+ipw, 0, nBlocksFm, ifhp, ifwp, 64);
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
              __m512i lcl_vnewmask   = _mm512_add_epi32( lcl_viadd, _mm512_set1_epi32((hi+kh)*ifw*16 + (wi+kw)*16) );
              __m512i lcl_vnewmask2  = _mm512_add_epi32( lcl_viadd, _mm512_set1_epi32((hi+kh)*ifw*16 + (wi+kw)*16) );
              __m512i lcl_vnewmask3  = _mm512_add_epi32( lcl_viadd, _mm512_set1_epi32((hi+kh)*ifw*16 + (wi+kw)*16) );
              __m512i lcl_vnewmask4  = _mm512_add_epi32( lcl_viadd, _mm512_set1_epi32((hi+kh)*ifw*16 + (wi+kw)*16) );
              __m512 lcl_vinput  = _mm512_load_act( input_ptr );
              __m512 lcl_vinput2 = _mm512_load_act( input_ptr+16 );
              __m512 lcl_vinput3 = _mm512_load_act( input_ptr+32 );
              __m512 lcl_vinput4 = _mm512_load_act( input_ptr+48 );
              __mmask16 lcl_mlt  = _mm512_cmp_ps_mask( lcl_voutput,  lcl_vinput,  _CMP_LT_OS );
              __mmask16 lcl_mlt2 = _mm512_cmp_ps_mask( lcl_voutput2, lcl_vinput2, _CMP_LT_OS );
              __mmask16 lcl_mlt3 = _mm512_cmp_ps_mask( lcl_voutput3, lcl_vinput3, _CMP_LT_OS );
              __mmask16 lcl_mlt4 = _mm512_cmp_ps_mask( lcl_voutput4, lcl_vinput4, _CMP_LT_OS );
              lcl_voutput   = _mm512_mask_blend_ps( lcl_mlt,  lcl_voutput,  lcl_vinput );
              lcl_voutput2  = _mm512_mask_blend_ps( lcl_mlt2, lcl_voutput2, lcl_vinput2 );
              lcl_voutput3  = _mm512_mask_blend_ps( lcl_mlt3, lcl_voutput3, lcl_vinput3 );
              lcl_voutput4  = _mm512_mask_blend_ps( lcl_mlt4, lcl_voutput4, lcl_vinput4 );
              lcl_vmask  = _mm512_mask_blend_epi32( lcl_mlt,  lcl_vmask,    lcl_vnewmask );
              lcl_vmask2 = _mm512_mask_blend_epi32( lcl_mlt2, lcl_vmask2,   lcl_vnewmask2 );
              lcl_vmask3 = _mm512_mask_blend_epi32( lcl_mlt3, lcl_vmask3,   lcl_vnewmask3 );
              lcl_vmask4 = _mm512_mask_blend_epi32( lcl_mlt4, lcl_vmask4,   lcl_vnewmask4 );
  #endif
  #if defined(LIBXSMM_DNN_POOLING_FWD_AVG)
              lcl_voutput  = _mm512_add_ps( lcl_voutput,  _mm512_load_act( input_ptr ) );
              lcl_voutput2 = _mm512_add_ps( lcl_voutput2, _mm512_load_act( input_ptr+16 ) );
              lcl_voutput3 = _mm512_add_ps( lcl_voutput3, _mm512_load_act( input_ptr+32 ) );
              lcl_voutput4 = _mm512_add_ps( lcl_voutput4, _mm512_load_act( input_ptr+48 ) );
  #endif
            }
          }
        }
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
        _mm512_storeu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw,  0, nBlocksFm, ofh, ofw, 64), lcl_vmask );
        _mm512_storeu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 16, nBlocksFm, ofh, ofw, 64), lcl_vmask2 );
        _mm512_storeu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 32, nBlocksFm, ofh, ofw, 64), lcl_vmask3 );
        _mm512_storeu_si512( &LIBXSMM_VLA_ACCESS(5, mask, img, fm, ho-oph, wo-opw, 48, nBlocksFm, ofh, ofw, 64), lcl_vmask4 );
  #endif
        _mm512_storeu_ps( lcl_output_ptr,    lcl_voutput );
        _mm512_storeu_ps( lcl_output_ptr+16, lcl_voutput2 );
        _mm512_storeu_ps( lcl_output_ptr+32, lcl_voutput3 );
        _mm512_storeu_ps( lcl_output_ptr+48, lcl_voutput4 );
      }
    }
  
    /* copy the local buffer into output activations */
    for( ho = oph; ho < (ofh+oph); ho++ ) {
      element_output_type*     output_ptr = &LIBXSMM_VLA_ACCESS(5, output,     img, fm,     ho, opw, 0, nBlocksFm, ofhp, ofwp, 64);
      float*               lcl_output_ptr = &LIBXSMM_VLA_ACCESS(3, lcl_output,          ho-oph,   0, 0,                   ofw, 64);
      for( wo = opw; wo < (ofw+opw); wo++ ) {
  #if defined(LIBXSMM_DNN_POOLING_FWD_AVG)
        const __m512 recp_pool_size_ps = _mm512_set1_ps( recp_pool_size );
        _mm512_stream_act( output_ptr,    _mm512_mul_ps( _mm512_loadu_ps( lcl_output_ptr ),    recp_pool_size_ps ) );
        _mm512_stream_act( output_ptr+16, _mm512_mul_ps( _mm512_loadu_ps( lcl_output_ptr+16 ), recp_pool_size_ps ) );
        _mm512_stream_act( output_ptr+32, _mm512_mul_ps( _mm512_loadu_ps( lcl_output_ptr+32 ), recp_pool_size_ps ) );
        _mm512_stream_act( output_ptr+48, _mm512_mul_ps( _mm512_loadu_ps( lcl_output_ptr+48 ), recp_pool_size_ps ) );
  #endif
  #if defined(LIBXSMM_DNN_POOLING_FWD_MAX)
        _mm512_stream_act( output_ptr,    _mm512_loadu_ps( lcl_output_ptr ) );
        _mm512_stream_act( output_ptr+16, _mm512_loadu_ps( lcl_output_ptr+16 ) );
        _mm512_stream_act( output_ptr+32, _mm512_loadu_ps( lcl_output_ptr+32 ) );
        _mm512_stream_act( output_ptr+48, _mm512_loadu_ps( lcl_output_ptr+48 ) );
  #endif
        output_ptr += 64;
        lcl_output_ptr += 64;
      }
    }
  }
  
  libxsmm_barrier_wait(handle->barrier, ltid);
  
  # undef _mm512_load_act
  # undef _mm512_stream_act
  # undef _mm512_store_act
  
  