bool gemm()

in src/c_interface.cpp [154:256]

• 80 lines of code
• 31 McCabe index (conditional complexity)

bool gemm(std::string precision, void *A, void *B, void *C,
          bool a_t, bool b_t,
          int m, int n, int k,
          int lda, int ldb, int ldc,
          float alpha, float beta,
          CUstream stream, unsigned int grid, unsigned int shared) {
  std::string kernel_op = get_op_string(a_t, b_t);

  if (grid >= selections[precision][kernel_op].size())
    return false;

  kernel_properties kp = selections[precision][kernel_op][grid];

  if (shared >= kp.shared_sizes.size())
    return false;

  bool vec4A, vec8A;
  bool vec4B, vec8B;
  if (a_t) {
    vec4A = (lda & 3) == 0 && (m & 3) == 0; //multiple of 4
    vec8A = (lda & 7) == 0 && (m & 7) == 0; //multiple of 8
  }
  else {
    vec4A = (lda & 3) == 0 && (k & 3) == 0; //multiple of 4
    vec8A = (lda & 7) == 0 && (k & 7) == 0; //multiple of 8
  }

  if (b_t) {
    vec4B = (ldb & 3) == 0 && (k & 3) == 0; //multiple of 4
    vec8B = (ldb & 7) == 0 && (k & 7) == 0; //multiple of 8
  }
  else {
    vec4B = (ldb & 3) == 0 && (n & 3) == 0; //multiple of 4
    vec8B = (ldb & 7) == 0 && (n & 7) == 0; //multiple of 8
  }

  bool vec4C = (ldc & 3) == 0 && (n & 3) == 0;

  bool vecA = (kp.vA == 4 && vec4A) || (kp.vA == 8 && vec8A);
  bool vecB = (kp.vB == 4 && vec4B) || (kp.vB == 8 && vec8B);
  bool vecC = kp.vC == 1 || vec4C;

  bool vec = vecA && vecB && vecC;

  CUdevice device;
  int res = cuCtxGetDevice(&device);
  if (res != 0)
    return false;

  bool success;
  int major, minor;

  std::tie(success, major, minor) = getDeviceProperties(device);

  if (success != 0)
    return false;

  std::string kernel_string;
  kernel_string.reserve(64);

  kernel_string += precision + "gemm_" + kp.tile_string +
                   "_" + kernel_op;
  if (vec)
    kernel_string += "_vec";

  if (major == 5)
    kernel_string += "_sm_50";
  else if (major == 6)
    kernel_string += "_sm_60";
  else
    return false;

  auto kernel = kernels_.find(kernel_string);
  if (kernel == kernels_.end()) {
    loadKernels(major);
    kernel = kernels_.find(kernel_string);
  }

  int blk_A = (m + kp.tile_m - 1) / kp.tile_m;
  int blk_B = (n + kp.tile_n - 1) / kp.tile_n;

  int blk_a, blk_b;
  std::tie(blk_a, blk_A) = closest_divisor(blk_A, kp.div);
  std::tie(blk_b, blk_B) = closest_divisor(blk_B, kp.div);

  if (blk_a == 1)
    std::tie(blk_a, blk_A) = std::make_pair(blk_A, 1);

  void *args[13] = {&C, &A, &B, &alpha, &beta, &lda, &ldb, &ldc,
                    &m, &n, &k, &blk_a, &blk_b};

  res = cuLaunchKernel(kernel->second, blk_a * blk_b, blk_B, blk_A,
                       kp.threads, 1, 1,
                       kp.shared_sizes[shared], stream, args, NULL);

  if (res != 0) {
    std::cerr << "Failed to execute " << kernel_string << " " <<
      res << std::endl;
    return false;
  }

  return true;
}