#include "ulp_neon.h"
#include "caffe2/core/timer.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

// TODO: tune this with cache size detection code. Changing to 32 helps on some
// devices (Snapdragon 820).
constexpr size_t kL1CacheSizeBytes = 16 * 1024;

#if defined(__ARM_NEON__) || defined(__ARM_NEON)

// Applies 2-bit uniform quantization to the floating point data at Xdata,
// storing QC bytes into XQdata (i.e. reading 8 * QC floats from Xdata).
// Requires QC to be a multiple of 8.
inline void quantize2bNeon(size_t QC,
                           const float* __restrict__ Xdata,
                           float offset,
                           float inter_center_distance,
                           std::array<uint8_t*, k2b1bXBits> XQdata) {
  TORCH_DCHECK_EQ(QC % 8, 0);
  const auto offset_plus_2_inter_center_distance = vdupq_n_f32(offset + 2 * inter_center_distance);
  const auto offset_plus_inter_center_distance = vdupq_n_f32(offset + inter_center_distance);
  const auto offset_ = vdupq_n_f32(offset);
  const uint8x8_t shifts = {1 << 0, 1 << 1, 1 << 2, 1 << 3, 1 << 4, 1 << 5, 1 << 6, 1 << 7};

  for (size_t qc = 0; qc < QC; qc += 8) {
    std::array<std::array<uint8x8_t, 8>, k2b1bXBits> ps;
    for (auto i = 0; i < k2b1bXBits; ++i) {
      for (auto j = 0; j < 8; ++j) {
        ps[i][j] = vdup_n_u8(0);
      }
    }

    for (auto j = 0; j < 8; ++j) {
      const auto x0 = vld1q_f32(&Xdata[qc * 8 + j * 8 + 0]);
      const auto x1 = vld1q_f32(&Xdata[qc * 8 + j * 8 + 4]);

      // logic.
      // if (v >= offset + inter_center_distance) {
      //   p[1] |= 1 << b;
      // } else {
      //   p[1] |= 0 << b;
      // }

      // if ((v >= offset && v < offset + inter_center_distance) ||
      //     (v >= offset * 2 * inter_center_distance)) {
      //   p[0] |= 1 << b;
      // } else {
      //   p[0] |= 0 << b;
      // }

      auto join = [](uint32x4_t a, uint32x4_t b) -> uint8x8_t {
        return vmovn_u16(vcombine_u16(vmovn_u32(a), vmovn_u32(b)));
      };

      const auto x_geq_offset_plus_2_inter_center_distance =
          join(vcgeq_s32(vreinterpretq_s32_f32(x0),
                         vreinterpretq_s32_f32(offset_plus_2_inter_center_distance)),
               vcgeq_s32(vreinterpretq_s32_f32(x1),
                         vreinterpretq_s32_f32(offset_plus_2_inter_center_distance)));
      const auto x_ge_offset =
          join(vcgeq_s32(vreinterpretq_s32_f32(x0), vreinterpretq_s32_f32(offset_)),
               vcgeq_s32(vreinterpretq_s32_f32(x1), vreinterpretq_s32_f32(offset_)));

      const auto x_lt_offset_plus_inter_center_distance =
          join(vcltq_s32(vreinterpretq_s32_f32(x0),
                         vreinterpretq_s32_f32(offset_plus_inter_center_distance)),
               vcltq_s32(vreinterpretq_s32_f32(x1),
                         vreinterpretq_s32_f32(offset_plus_inter_center_distance)));

      const auto p1_mask = vmvn_u8(x_lt_offset_plus_inter_center_distance);
      const auto p0_mask = vorr_u8(vand_u8(x_ge_offset, x_lt_offset_plus_inter_center_distance),
                                   x_geq_offset_plus_2_inter_center_distance);
      ps[0][j] = vand_u8(shifts, p0_mask);
      ps[1][j] = vand_u8(shifts, p1_mask);
    }

    for (auto i = 0; i < 2; ++i) {
      const auto p01 = vpadd_u8(ps[i][0], ps[i][1]);
      const auto p23 = vpadd_u8(ps[i][2], ps[i][3]);
      const auto p45 = vpadd_u8(ps[i][4], ps[i][5]);
      const auto p67 = vpadd_u8(ps[i][6], ps[i][7]);
      const auto p0123 = vpadd_u8(p01, p23);
      const auto p4567 = vpadd_u8(p45, p67);
      vst1_u8(XQdata[i] + qc, vpadd_u8(p0123, p4567));
    }
  }
}

void uniformQuantize2b1bNeon(QConvState* state,
                             const TensorCPU& X,
                             const std::vector<std::unique_ptr<TensorCPU>>& XQ,
                             float offset,
                             float inter_center_distance) {
  CAFFE_ENFORCE_GT(X.ndim(), 1);
  const size_t C = X.dim32(X.ndim() - 1);
  const size_t N = X.size() / C;
  const size_t QC = divRoundUp(C, 8);
  auto XQs = X.sizes().vec();
  XQs[X.ndim() - 1] = QC;
  CAFFE_ENFORCE_EQ(XQ.size(), k2b1bXBits);
  for (auto i = 0; i < k2b1bXBits; ++i) {
    XQ[i]->Resize(XQs);
  }
  const float* Xdata = X.data<float>();
  std::array<uint8_t*, k2b1bXBits> XQdata;
  for (size_t i = 0; i < k2b1bXBits; ++i) {
    XQdata[i] = XQ[i]->mutable_data<uint8_t>();
  }
  CAFFE_ENFORCE_GT(offset, 0);
  CAFFE_ENFORCE_GT(inter_center_distance, 0);
  size_t QCUnroll = ((C / 8) / 8) * 8;
  // Each worker loads an L1 cache sized block.
  // We read/write B * K * 4 + 2 * B * (K / 8), so to fit inside C, we have
  // B = 4 * C / 17 K.
  // QCUnroll = 0;
  const size_t rowsPerBlock =
      std::max<size_t>(std::floor<size_t>(double(4 * kL1CacheSizeBytes) / double(17 * C)), 1);
  state->parallelFor(divRoundUp(N, rowsPerBlock), [&](size_t nb) {
    for (size_t n = nb * rowsPerBlock; n < std::min<size_t>(nb * rowsPerBlock + rowsPerBlock, N);
         ++n) {
      std::array<uint8_t*, k2b1bXBits> XQoff = {{
          XQdata[0] + 0 + QC * n, XQdata[1] + 0 + QC * n,
      }};
      quantize2bNeon(QCUnroll, &Xdata[0 + C * n], offset, inter_center_distance, XQoff);
      for (size_t qc = QCUnroll; qc < QC; ++qc) {
        // compute the block in X.
        std::array<uint8_t, k2b1bXBits> p = {{0, 0}};
        for (size_t b = 0; b < 8; ++b) {
          const size_t c = qc * 8 + b;
          if (c < C) {
            float v = Xdata[c + C * n];
            if (v < offset) {
              // zero'd already.
            } else if (v < offset + inter_center_distance) {
              p[0] |= 1 << b;
            } else if (v < offset + 2 * inter_center_distance) {
              p[1] |= 1 << b;
            } else {
              p[0] |= 1 << b;
              p[1] |= 1 << b;
            }
          }
        }
        for (auto i = 0; i < k2b1bXBits; ++i) {
          XQdata[i][qc + QC * n] = p[i];
        }
      }
    }
  });
}

template <size_t TileSize, size_t TileDepthBytes>
void uniformQuantize2b1bNeonPacked(QConvState* state,
                                   const TensorCPU& X,
                                   const std::vector<std::unique_ptr<TensorCPU>>& XQ,
                                   float offset,
                                   float inter_center_distance) {
  const size_t M = X.size_to_dim(3);
  const size_t K = X.size() / M;
  const size_t QK = divRoundUp(K, 8);
  const size_t numTiles = divRoundUp(M, TileSize);
  const size_t numTilesDepth = divRoundUp(QK, TileDepthBytes);
  for (size_t i = 0; i < k2b1bXBits; ++i) {
    XQ[i]->Resize(numTiles, numTilesDepth, TileSize, TileDepthBytes);
  }
  const float* Xdata = X.data<float>();
  std::array<uint8_t*, k2b1bXBits> XQdata;
  for (auto i = 0; i < k2b1bXBits; ++i) {
    XQdata[i] = XQ[i]->mutable_data<uint8_t>();
  }
  CAFFE_ENFORCE_GT(offset, 0);
  CAFFE_ENFORCE_GT(inter_center_distance, 0);
  // Each worker loads an L1 cache sized block.
  // We read/write B * K * TileSize * 4 + 2 * B * TileSize * (K / 8), so to fit inside C, we have
  // B = 4 * C / (17 * K * TileSize).
  const size_t tilesPerBlock = std::max<size_t>(
      std::floor<size_t>(double(4 * kL1CacheSizeBytes) / double(17 * K * TileSize)), 1);
  state->parallelFor(divRoundUp(numTiles, tilesPerBlock), [&](size_t nb) {
    for (size_t i = nb * tilesPerBlock;
         i < std::min<size_t>(nb * tilesPerBlock + tilesPerBlock, numTiles);
         ++i) {
      for (size_t j = 0; j < numTilesDepth; ++j) {
        if (i != numTiles - 1 && j != numTilesDepth - 1) {
          // we have a full tile. Just memcpy.
          for (auto ii = 0; ii < TileSize; ++ii) {
            size_t m = i * TileSize + ii;
            size_t k = j * TileDepthBytes * 8;
            std::array<uint8_t*, k2b1bXBits> XQoff = {
                {XQdata[0] + TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                     TileSize * TileDepthBytes * numTilesDepth * i,
                 XQdata[1] + TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                     TileSize * TileDepthBytes * numTilesDepth * i}};
            quantize2bNeon(TileDepthBytes, &Xdata[m * K + k], offset, inter_center_distance, XQoff);
          }
        } else {
          for (size_t ii = 0; ii < TileSize; ++ii) {
            size_t m = i * TileSize + ii;
            size_t k = j * TileDepthBytes * 8;
            std::array<uint8_t*, k2b1bXBits> XQoff = {
                {XQdata[0] + TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                     TileSize * TileDepthBytes * numTilesDepth * i,
                 XQdata[1] + TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                     TileSize * TileDepthBytes * numTilesDepth * i}};
            if (m < M && k + TileDepthBytes * 8 <= K) {
              // We can just read the stripe directly.
              quantize2bNeon(
                  TileDepthBytes, &Xdata[m * K + k], offset, inter_center_distance, XQoff);
            } else {
              // We need to pad the stripe to the full amount read by
              // quantize2bNeon.
              std::array<float, 8 * TileDepthBytes> Xpad = {{0}};
              if (m < M) {
                std::copy(&Xdata[m * K + k], &Xdata[m * K + K], Xpad.begin());
              }
              quantize2bNeon(TileDepthBytes, Xpad.data(), offset, inter_center_distance, XQoff);
            }
          }
        }
      }
    }
  });
}

// Packs a matrix (of size MxK) into a tiled array of size
// (M/TileSize)x(K/TileDepthBytes)xTileSizexTileDepthBytes.
template <size_t TileSize, size_t TileDepthBytes>
void qpack_tiles(QConvState* state, const TensorCPU& X, size_t axis, TensorCPU* XP) {
  const size_t M = X.size_to_dim(axis);
  const size_t QK = X.size() / M;
  const size_t numTiles = divRoundUp(M, TileSize);
  const size_t numTilesDepth = divRoundUp(QK, TileDepthBytes);
  XP->Resize(numTiles, numTilesDepth, TileSize, TileDepthBytes);

  const auto* __restrict__ Xdata = X.data<uint8_t>();
  auto* __restrict__ XPdata = XP->mutable_data<uint8_t>();
  // Load L1 sized tiles per thread.
  // We read/write 2 * B * QK * TileSize bytes, so
  // B = C / (2 * QK * TileSize)
  const size_t tilesPerBlock = std::max<size_t>(
      std::floor<size_t>(double(kL1CacheSizeBytes) / double(2 * TileSize * QK)), 1);
  state->parallelFor(divRoundUp(numTiles, tilesPerBlock), [&](size_t nb) {
    for (size_t i = nb * tilesPerBlock;
         i < std::min<size_t>(nb * tilesPerBlock + tilesPerBlock, numTiles);
         ++i) {
      for (size_t j = 0; j < numTilesDepth; ++j) {
        if (i != numTiles - 1 && j != numTilesDepth - 1) {
          // we have a full tile. Just memcpy.
          for (auto ii = 0; ii < TileSize; ++ii) {
            auto m = i * TileSize + ii;
            auto qk = j * TileDepthBytes;
            std::memcpy(&XPdata[TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                                TileSize * TileDepthBytes * numTilesDepth * i],
                        &Xdata[m * QK + qk],
                        TileDepthBytes);
          }
        } else {
          for (size_t ii = 0; ii < TileSize; ++ii) {
            for (size_t jj = 0; jj < TileDepthBytes; ++jj) {
              size_t m = i * TileSize + ii;
              size_t qk = j * TileDepthBytes + jj;
              uint8_t pval = 0;
              if (m < M && qk < QK) {
                // get value from X
                pval = Xdata[m * QK + qk];
              }
              XPdata[jj + TileDepthBytes * ii + TileDepthBytes * TileSize * j +
                     TileSize * TileDepthBytes * numTilesDepth * i] = pval;
            }
          }
        }
      }
    }
  });
}

// Computes the kUnrollM x kUnrollM tile of a GEMM by multiplying two packed
// slices of size (kUnrolLMxK). These tiles are constructed by the qpack_tiles
// function, which packs an input array of size [M][K] into an
// [M/TileSize][K/TileDepthBytes][TileSize][TileDepthBytes], which ensures all
// the array accesses in this function is contiguous.
template <size_t kUnrollM, size_t kUnrollN, size_t TileDepthBytes, typename F>
void qgess_packed(const uint8_t* __restrict__ Ablock,
                  const uint8_t* __restrict__ Bblock,
                  float* __restrict__ Cblock,
                  const size_t Cstride,
                  const size_t QK,
                  const size_t Nstart,
                  F&& f) {
  static_assert(kUnrollN % 8 == 0, "");
  static_assert(TileDepthBytes == 16, "");
  TORCH_DCHECK_EQ(QK % 16, 0);
  uint16x8_t acc[kUnrollM][kUnrollN / 8];
  for (size_t mm = 0; mm < kUnrollM; ++mm) {
    for (size_t nn = 0; nn < kUnrollN / 8; ++nn) {
      acc[mm][nn] = vdupq_n_u16(0);
    }
  }
  size_t qk = 0;
  const size_t QK16Unroll = (QK / 16) * 16;
  for (; qk < QK16Unroll; qk += 16) {
    uint8x16_t Areg[kUnrollM];
    for (size_t mm = 0; mm < kUnrollM; ++mm) {
      Areg[mm] = vld1q_u8(Ablock);
      Ablock += 16;
    }

    for (size_t nn = 0; nn < kUnrollN / 8; ++nn) {
      uint8x16_t Breg[8];
      for (size_t nnn = 0; nnn < 8; ++nnn) {
        Breg[nnn] = vld1q_u8(Bblock);
        Bblock += 16;
      }
      for (size_t mm = 0; mm < kUnrollM; ++mm) {
        uint8x16_t cnts[8];
        for (size_t nnn = 0; nnn < 8; ++nnn) {
          cnts[nnn] = vcntq_u8(veorq_u8(Breg[nnn], Areg[mm]));
        }
        uint8x8_t ps[8];
        for (size_t nnn = 0; nnn < 8; ++nnn) {
          ps[nnn] = vadd_u8(vget_low_u8(cnts[nnn]), vget_high_u8(cnts[nnn]));
        }
        uint8x8_t pss[4];
        for (size_t nnn = 0; nnn < 4; ++nnn) {
          pss[nnn] = vpadd_u8(ps[2 * nnn], ps[2 * nnn + 1]);
        }
        uint8x8_t psss[2];
        for (size_t nnn = 0; nnn < 2; ++nnn) {
          psss[nnn] = vpadd_u8(pss[2 * nnn], pss[2 * nnn + 1]);
        }
        uint8x16_t out = vcombine_u8(psss[0], psss[1]);
        acc[mm][nn] = vpadalq_u8(acc[mm][nn], out);
      }
    }
  }

  for (size_t mm = 0; mm < kUnrollM; ++mm) {
    auto* Crow = Cblock + mm * Cstride;
    for (size_t nn = 0; nn < kUnrollN / 8; ++nn) {
      const int32x4_t K_ = vdupq_n_s32(QK * 8);
      const int16x4_t two = vdup_n_s16(2);
      const int16x4_t acc0123_l = vreinterpret_s16_u16(vget_low_u16(acc[mm][nn]));
      const int16x4_t acc0123_h = vreinterpret_s16_u16(vget_high_u16(acc[mm][nn]));
      const int32x4_t K_minus_2_acc0123_l = vmlsl_s16(K_, two, acc0123_l);
      const int32x4_t K_minus_2_acc0123_h = vmlsl_s16(K_, two, acc0123_h);
      f(Crow + nn * 8 + 0, vcvtq_f32_s32(K_minus_2_acc0123_l), Nstart + nn * 8 + 0);
      f(Crow + nn * 8 + 4, vcvtq_f32_s32(K_minus_2_acc0123_h), Nstart + nn * 8 + 4);
    }
  }
}

// Computes the (normal + transpose) matrix-matrix product of two -1/1 binary
// matrices, laid out in the standard format.
template <size_t TileSize, size_t TileDepthBytes, typename F>
inline void qgemm_nt_packed(
    QConvState* state, const TensorCPU& A, const TensorCPU& B, TensorCPU* C, F&& f = F()) {
  CAFFE_ENFORCE_EQ(A.ndim(), 4);
  CAFFE_ENFORCE_EQ(B.ndim(), 4);
  CAFFE_ENFORCE_EQ(A.dim(2), TileSize);
  CAFFE_ENFORCE_EQ(B.dim(2), TileSize);
  CAFFE_ENFORCE_EQ(A.dim(3), TileDepthBytes);
  CAFFE_ENFORCE_EQ(B.dim(3), TileDepthBytes);
  const size_t MT = A.dim(0);
  const size_t NT = B.dim(0);
  const size_t M = MT * TileSize;
  const size_t N = NT * TileSize;

  const size_t QKT = A.dim(1);
  const size_t K = QKT * 8 * TileDepthBytes;
  const size_t QK = K / 8;
  CAFFE_ENFORCE_EQ(A.dim(1), B.dim(1));
  C->Resize(M, N);
  const auto* Adata = A.data<uint8_t>();
  const auto* Bdata = B.data<uint8_t>();
  auto* Cdata = C->mutable_data<float>();

  // Assume TxT tile. Each input slice is of size T x (K/8) bytes, and the output
  // is a tile of size T x T x sizeof(float) bytes. We want the sum of this to fit
  // in L1 cache. This means for a block number of tiles B , we load B * T * K /
  // 8 + B * T * K / 8 + B * B * T * T * sizeof(float).

  // If cache size = C, we get
  // B = 1/(32 * T) (sqrt(256 C + K^2) - K)
  // taking floor (by integer division), gives the result.

  // Assume 16KB L1 cache.
  size_t tilesPerBlock =
      std::floor((std::sqrt(256 * kL1CacheSizeBytes + K * K) - K) / (32 * TileSize));
  if (tilesPerBlock < 1) {
    tilesPerBlock = 1;
  }
  CAFFE_ENFORCE_LT(K, std::pow(2, 16));
  CAFFE_ENFORCE_EQ(M % TileSize, 0);
  CAFFE_ENFORCE_EQ(N % TileSize, 0);
  const size_t MNumTiles = M / TileSize;
  const size_t NNumTiles = N / TileSize;
  const size_t MNumBlocks = divRoundUp(MNumTiles, tilesPerBlock);
  const size_t NNumBlocks = divRoundUp(NNumTiles, tilesPerBlock);

  state->parallelFor(MNumBlocks * NNumBlocks, [&](size_t mn) {
    const size_t mBlockIdx = mn / NNumBlocks;
    const size_t nBlockIdx = mn % NNumBlocks;
    const size_t mTileStart = mBlockIdx * tilesPerBlock;
    const size_t nTileStart = nBlockIdx * tilesPerBlock;
    for (size_t mBlockTileIdx = 0;
         mBlockTileIdx < tilesPerBlock && mBlockTileIdx + mTileStart < MNumTiles;
         ++mBlockTileIdx) {
      const size_t mTileIdx = mBlockTileIdx + mTileStart;
      for (size_t nBlockTileIdx = 0;
           nBlockTileIdx < tilesPerBlock && nBlockTileIdx + nTileStart < NNumTiles;
           ++nBlockTileIdx) {
        const size_t nTileIdx = nBlockTileIdx + nTileStart;
        // A layout: [M/TileSize][QK / TileDepth][TileSize][TileDepth]
        // C layout: [M/TileSize][TileSize][N/TileSize][TileSize]
        const auto* Ablock = &Adata[mTileIdx * QK * TileSize];
        const auto* Bblock = &Bdata[nTileIdx * QK * TileSize];
        auto* Cblock = &Cdata[mTileIdx * TileSize * N + nTileIdx * TileSize];
        const size_t Cstride = N;
        qgess_packed<TileSize, TileSize, TileDepthBytes, F>(
            Ablock, Bblock, Cblock, Cstride, QK, nTileIdx * TileSize, std::forward<F>(f));
      }
    }
  });
}

void run2b1bConvIm2ColGEMM(QConvState* state,
                           const ConvArgs& args,
                           const TensorCPU& X,
                           TensorCPU* Y) {
  // TODO: packing + quantization in same block.
  const size_t KH = state->WQ->dim32(1);
  const size_t KW = state->WQ->dim32(2);
  const size_t OH = (X.dim32(1) - KH + args.pad_t + args.pad_b) / args.stride_h + 1;
  const size_t OW = (X.dim32(2) - KW + args.pad_l + args.pad_r) / args.stride_w + 1;
  const size_t OC = state->WQ->dim32(0);
  const size_t QK = KH * KW * divRoundUp(X.dim32(3), 8);
  Y->Resize(X.dim32(0), OH, OW, OC);
  if (!state->WQPacked) {
    state->WQPacked = std::make_unique<Tensor>(CPU);
    qpack_tiles<kGEMMTileSize, kGEMMTileDepthBytes>(state, *(state->WQ), 1, state->WQPacked.get());
    CAFFE_ENFORCE_EQ(state->WQPacked->dim32(0), divRoundUp(OC, kGEMMTileSize));
    CAFFE_ENFORCE_EQ(state->WQPacked->dim32(1), divRoundUp(QK, kGEMMTileDepthBytes));
    CAFFE_ENFORCE_EQ(state->WQPacked->dim32(2), kGEMMTileSize);
    CAFFE_ENFORCE_EQ(state->WQPacked->dim32(3), kGEMMTileDepthBytes);

    // We can fuse the bias addition into the filter normalization. We can
    // replace the bias + 3/2 normalization factor by replacing normalization
    // with (2/3 bias + normalization), and setting bias to zero.
    if (state->bias) {
      for (size_t i = 0; i < state->bias->size(); ++i) {
        state->WQN->mutable_data<float>()[i] += 2.0f / 3 * state->bias->data<float>()[i];
      }
    }
    state->bias.reset();

    // If we have to pad when we pack our weight tiles, then we need to adjust
    // the normalization factor by the number of zeros that we added.
    const size_t QKPadding = divRoundUp(QK, kGEMMTileDepthBytes) * kGEMMTileDepthBytes - QK;
    if (QKPadding != 0) {
      for (size_t i = 0; i < state->WQN->size(); ++i) {
        state->WQN->mutable_data<float>()[i] -= QKPadding * 8;
      }
    }
  }
  CAFFE_ENFORCE(!state->bias.get());
  // Since 1x1s are so common, we fuse the quantization + packing steps.
  const bool is_1x1 = KH == 1 && KW == 1 && args.pad_l == 0 && args.pad_r == 0 && args.pad_b == 0 &&
                      args.pad_t == 0 && args.stride_h == 1 && args.stride_w == 1;

  if (is_1x1) {
    CAFFE_ENFORCE_EQ(OH, X.dim32(1));
    CAFFE_ENFORCE_EQ(OW, X.dim32(2));
    uniformQuantize2b1bNeonPacked<kGEMMTileSize, kGEMMTileDepthBytes>(
        state, X, state->XQs, 0.5, 1.0);
  } else {
    uniformQuantize2b1bNeon(state, X, state->XQs, 0.5, 1.0);
  }
  TensorCPU* YQ0 = state->YQs[0].get();

  if (state->WQ->dim32(0) % kGEMMTileSize == 0) {
    // We can run inplace by operating on our Y vector, and then shrinking Y.
    YQ0 = Y;
  }

  for (size_t i = 0; i < k2b1bXBits; ++i) {
    const auto& XQ = *(state->XQs[i]);
    if (!is_1x1) {
      qim2col(args, XQ, *(state->WQ), state->scratchColBuffer.get());
      qpack_tiles<kGEMMTileSize, kGEMMTileDepthBytes>(
          state, *(state->scratchColBuffer), 3, state->scratch.get());
    }

    {
      const auto* __restrict__ WQNdata = state->WQN->data<float>();
      switch (i) {
      case 0:
        qgemm_nt_packed<kGEMMTileSize, kGEMMTileDepthBytes>(
            state,
            is_1x1 ? XQ : *(state->scratch),
            *(state->WQPacked),
            YQ0,
            [WQNdata](float* __restrict__ acc, float32x4_t value, size_t channel) {
              // acc[c] = 3/2 WQN[c] + 1/2 value[c];
              const float32x4_t _32 = vdupq_n_f32(3.0f / 2);
              const float32x4_t _12 = vdupq_n_f32(1.0f / 2);
              const float32x4_t WQNc_32 = vmulq_f32(_32, vld1q_f32(WQNdata + channel));
              const float32x4_t WQNc_32_value_12 = vmlaq_f32(WQNc_32, _12, value);
              vst1q_f32(acc, WQNc_32_value_12);
            });
        break;
      case 1:
        qgemm_nt_packed<kGEMMTileSize, kGEMMTileDepthBytes>(
            state,
            is_1x1 ? XQ : *(state->scratch),
            *(state->WQPacked),
            YQ0,
            [](float* __restrict__ acc, float32x4_t value, size_t channel) {
              const float32x4_t curr = vld1q_f32(acc);
              vst1q_f32(acc, vaddq_f32(curr, value));
            });
        break;
      }
    }
  }

  if (YQ0 != Y) {
    // In this case, the stride does not match, so we need to copy the output
    // data into the contiguous Y matrix.
    const size_t F = state->WQ->dim(0);
    const size_t N = Y->size() / F;
    const size_t NP = YQ0->dim32(0);
    const size_t FP = YQ0->dim32(1);
    math::CopyMatrix<CPUContext>(
        sizeof(float), N, F, YQ0->data<float>(), FP, Y->mutable_data<float>(), F, nullptr);
  } else {
    CAFFE_ENFORCE_EQ(Y->dim32(0), divRoundUp(X.dim32(0) * OH * OW, kGEMMTileSize) * kGEMMTileSize);
    CAFFE_ENFORCE_EQ(Y->dim32(1), OC);
    Y->ShrinkTo(X.dim32(0) * OH * OW);
    Y->Reshape(std::vector<int64_t>{{int64_t(X.dim(0)), int64_t(OH), int64_t(OW), int64_t(OC)}});
  }
}

bool run2b1bConvNeon(QConvState* state, const ConvArgs& args, const TensorCPU& X, TensorCPU* Y) {
  // TODO: insert specialized cases (e.g. depthwise convolutions, the direct
  // convolution.
  CAFFE_ENFORCE_EQ(X.ndim(), 4);
  run2b1bConvIm2ColGEMM(state, args, X, Y);
  return true;
}

#endif

} // namespace caffe2