void Compute_XN()

in src/bst_op.cc [251:320]

• 55 lines of code
• 13 McCabe index (conditional complexity)

  void Compute_XN(OpKernelContext* ctx, uint op, uint max_lut)
  {
    const Tensor& a   = ctx->input(0);
    const Tensor& b   = ctx->input(1);
    const Tensor& lut = ctx->input(2 + op);

    OP_REQUIRES(ctx, a.dims() == 5 && b.dims() == 3, errors::InvalidArgument("Mismatched Shapes: a,b"));
    OP_REQUIRES(ctx, lut.dims() == 3,                errors::InvalidArgument("Bad lut"));

    uint lut_heads = lut.dim_size(0);
    uint lut_dim   = lut.dim_size(1);
    uint batch_dim = b.dim_size(0);
    uint state_dim = b.dim_size(2);

    if (head_state_ == 0)
    {
      OP_REQUIRES(ctx,
        a.dim_size(0) == batch_dim &&
        a.dim_size(1) == heads_    &&
        a.dim_size(2) == blocks_   &&
        a.dim_size(3) == blk_size_ &&
        a.dim_size(4) == blk_size_, errors::InvalidArgument("Mismatched A shape"));

      head_state_ = state_dim / heads_;
      OP_REQUIRES(ctx, state_dim % heads_ == 0,                  errors::InvalidArgument("state_dim not evenly divisible by number of heads"));
      OP_REQUIRES(ctx, (head_state_ & 7) == 0,                   errors::InvalidArgument("Head state dim must be multiple of 8, and ideally a multiple of 64"));
      OP_REQUIRES(ctx, b.dim_size(1) == ctx_blks_b_ * blk_size_, errors::InvalidArgument("Bad B context length"));
      OP_REQUIRES(ctx, lut_heads == heads_ || lut_heads == 1,    errors::InvalidArgument("Bad head dim"));

      uint div = CEIL_DIV(head_state_, 64);
      magicu64(div, magic_, shift_);
      OP_REQUIRES(ctx, magic_ > 0, errors::InvalidArgument("Bad magic for div: ", div));
    }

    Tensor* c;
    TensorShape shapeC({batch_dim, ctx_blks_c_ * blk_size_, state_dim});
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, shapeC, &c));

    const uint2* l_ptr = (const uint2*)lut.flat<int32>().data();

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

    Benchmark* bench = nullptr;
    if (bench_) bench = new Benchmark(stream, bench_string_, 0, flops_ * (float)(batch_dim * state_dim), repeat_);

    if (a.dtype() == DT_HALF)
    {
      OP_REQUIRES(ctx, major_ >= 7, errors::InvalidArgument("Tensorcore GPU required"));

      const ehalf* a_ptr = (const ehalf*)a.tensor_data().data();
      const ehalf* b_ptr = (const ehalf*)b.tensor_data().data();
            ehalf* c_ptr = (      ehalf*)c->tensor_data().data();

      for (int r = 0; r < repeat_; r++)
        bst_hgemm_xn(stream, l_ptr, a_ptr, b_ptr, c_ptr, blk_size_, blocks_, batch_dim, ctx_blks_b_, ctx_blks_c_, heads_, head_state_, lut_heads, lut_dim, op, magic_, shift_, max_lut);
    }
    else
    {
      const bhalf* a_ptr = (const bhalf*)a.tensor_data().data();
      const float* b_ptr = (const float*)b.tensor_data().data();
            float* c_ptr = (      float*)c->tensor_data().data();

      OP_REQUIRES(ctx, blk_size_ == 32, errors::InvalidArgument("Only blocksize=32 supported for fp32 pathway."));

      for (int r = 0; r < repeat_; r++)
        bst_sgemm_xn(stream, l_ptr, a_ptr, b_ptr, c_ptr, blk_size_, blocks_, batch_dim, ctx_blks_b_, ctx_blks_c_, heads_, head_state_, lut_heads, lut_dim, op, magic_, shift_, max_lut);
    }

    if (bench) delete bench;
  }