virtual void Compute()

in src/blocksparse_conv_op.cc [267:372]
94 lines of code
14 McCabe index (conditional complexity)

  virtual void Compute(OpKernelContext* ctx) override {

    const Tensor& grid_lut = ctx->input(mode_);
    const Tensor& mpq_lut  = ctx->input(mode_ == 1 ? 4 : 3);
    const Tensor& ck_lut   = ctx->input(5);
    const Tensor& a        = ctx->input(6);
    const Tensor& b        = ctx->input(7);

    float alpha = 1.0f;
    int   gridX = grid_lut.dim_size(0);
    int   rank  = b.dims();
    int   N     = b.dim_size(0);

    int zero, gridY;
    TensorShape c_shape;
    if (mode_ == 0)
    {
      zero  = N * zero_;
      gridY = N;
      c_shape.AddDim(N);
      c_shape.AddDim(K_);
      if (rank == 5) c_shape.AddDim(MPQ_[0]);
      if (rank >= 4) c_shape.AddDim(MPQ_[1]);
                     c_shape.AddDim(MPQ_[2]);
    }
    else if (mode_ == 1)
    {
      zero  = N * zero_;
      gridY = N;
      c_shape.AddDim(N);
      c_shape.AddDim(C_);
      if (rank == 5) c_shape.AddDim(DHW_[0]);
      if (rank >= 4) c_shape.AddDim(DHW_[1]);
                     c_shape.AddDim(DHW_[2]);
    }
    else
    {
      zero  = zero_;
      gridY = 1;
      for (std::vector<int32>::iterator it = dimF_.begin() ; it != dimF_.end(); ++it)
        c_shape.AddDim(*it);
    }

    Tensor* c = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, c_shape, &c));

    CUdeviceptr    c_ptr = (CUdeviceptr)c->flat<CT>().data();
    CUdeviceptr grid_ptr = (CUdeviceptr)grid_lut.flat<int32>().data();
    CUdeviceptr  mpq_ptr = (CUdeviceptr)mpq_lut.flat<int32>().data();
    CUdeviceptr   ck_ptr = (CUdeviceptr)ck_lut.flat<int32>().data();
    CUdeviceptr    a_ptr = (CUdeviceptr)a.flat<AT>().data();
    CUdeviceptr    b_ptr = (CUdeviceptr)b.flat<BT>().data();

    OP_REQUIRES_OK(ctx, GetKernel(kernel_name_, &kernel_));

    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    void *args[] = {
      &grid_ptr, &mpq_ptr, &ck_ptr, &c_ptr, &a_ptr, &b_ptr, &alpha,
      &trs_, &magic_trs_, &shift_trs_, &cdhw_, &kmpq_, &N, &size_f_
    };

    CUresult res;
    if (zero > 0)
    {
      res = cuMemsetD8Async(c_ptr, 0, zero, stream);
      if (res != CUDA_SUCCESS)
      {
        const char* errstr;
        cuGetErrorString(res, &errstr);
        OP_REQUIRES(ctx, false, errors::Internal("cuMemsetD8Async Error: ", errstr, " bytes: ", zero));
      }
    }
    res = cuLaunchKernel(kernel_, gridX, gridY, 1, threads_, 1, 1, share_, stream, args, NULL);
    if (res != CUDA_SUCCESS)
    {
      const char* errstr;
      cuGetErrorString(res, &errstr);
      char params[256];
      sprintf(params, "m:%d(%5d,%5d:%5d), grid:%p, mpq:%p, ck:%p, c:%p, a:%p, b:%p, C:%5d, K:%5d, N:%3d %s\n",
        mode_, gridX, gridY, share_,
        (void*)grid_ptr, (void*)mpq_ptr, (void*)ck_ptr,
        (void*)c_ptr, (void*)a_ptr, (void*)b_ptr,
        C_, K_, N, kernel_name_.c_str());
      OP_REQUIRES(ctx, false, errors::Internal("cuLaunchKernel Error: ", errstr, "\nParams: ", params ));
      //OP_REQUIRES(ctx, false, errors::Internal("cuLaunchKernel Error: ", errstr));
    }
    if (debug_)
    {
      res = cuStreamSynchronize(stream);
      if (res != CUDA_SUCCESS)
      {
        const char* errstr;
        cuGetErrorString(res, &errstr);

        char params[256];
        sprintf(params, "m:%d(%5d,%5d:%5d), grid:%p, mpq:%p, ck:%p, c:%p, a:%p, b:%p, C:%5d, K:%5d, N:%3d %s\n",
          mode_, gridX, gridY, share_,
          (void*)grid_ptr, (void*)mpq_ptr, (void*)ck_ptr,
          (void*)c_ptr, (void*)a_ptr, (void*)b_ptr,
          C_, K_, N, kernel_name_.c_str());

          OP_REQUIRES(ctx, false, errors::Internal("Cuda Error: ", errstr, "\nParams: ", params ));
      }
    }
  }