void Compute_NT()

in src/bst_op.cc [183:250]

• 53 lines of code
• 11 McCabe index (conditional complexity)

  void Compute_NT(OpKernelContext* ctx)
  {
    const Tensor& a   = ctx->input(0);
    const Tensor& b   = ctx->input(1);
    const Tensor& lut = ctx->input(2);

    OP_REQUIRES(ctx, a.dims() == 3 && 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 = a.dim_size(0);
    uint state_dim = a.dim_size(2);

    if (head_state_ == 0)
    {
      OP_REQUIRES(ctx,
        b.dim_size(0) == batch_dim &&
        b.dim_size(2) == state_dim, errors::InvalidArgument("Mismatched Shapes"));

      OP_REQUIRES(ctx, a.dim_size(1) == ctx_blks_a_ * blk_size_, errors::InvalidArgument("Bad A context length"));
      OP_REQUIRES(ctx, b.dim_size(1) == ctx_blks_b_ * blk_size_, errors::InvalidArgument("Bad B context length"));

      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, lut_heads == heads_ || lut_heads == 1, errors::InvalidArgument("Bad head dim"));
    }

    Tensor* c;
    TensorShape shapeC({batch_dim, heads_, blocks_, blk_size_, blk_size_});
    OP_REQUIRES(ctx, shapeC.num_elements() < (1ull << 32), errors::InvalidArgument("Attention space too large.  Only 32 bit address offsets currently supported (easily fixed if needed.)"));
    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();

      for (int r = 0; r < repeat_; r++)
        if (c->dtype() == DT_HALF)
          bst_hgemm_nt<ehalf,ehalf2,ehalf4>(stream, l_ptr, a_ptr, b_ptr, (ehalf*)c->tensor_data().data(), blk_size_, blocks_, batch_dim, ctx_blks_a_, ctx_blks_b_, heads_, head_state_, lut_heads, lut_dim);
        else
          bst_hgemm_nt<bhalf,bhalf2,bhalf4>(stream, l_ptr, a_ptr, b_ptr, (bhalf*)c->tensor_data().data(), blk_size_, blocks_, batch_dim, ctx_blks_a_, ctx_blks_b_, heads_, head_state_, lut_heads, lut_dim);
    }
    else
    {
      const float* a_ptr = (const float*)a.tensor_data().data();
      const float* b_ptr = (const float*)b.tensor_data().data();
            bhalf* c_ptr = (bhalf*)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_nt(stream, l_ptr, a_ptr, b_ptr, c_ptr, blk_size_, blocks_, batch_dim, ctx_blks_a_, ctx_blks_b_, heads_, head_state_, lut_heads, lut_dim);
    }

    if (bench) delete bench;
  }