void Compute()

in src/transformer_op.cc [311:362]
36 lines of code
11 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override {

    const Tensor& dy = ctx->input(0);
    const Tensor&  y = ctx->input(1);
    const Tensor&  s = ctx->input(2);
    OpInputList m;     ctx->input_list("mask", &m);

    // y: D0, D1, D2, D3
    // m:  1, D1, D2, D3
    // m:  1,  1, D2, D3
    // m:  1,  1,  1, D3
    int rank = y.dims();
    uint D3 = y.dim_size(--rank);
    uint D2 = 1, D1 = 1, D0 = 1, M2 = 0, M1 = 0;
    const float* m_ptr = NULL;
    if (m.size() > 0)
    {
      // gather inner dimensions and strides of the mask
      // only dims 1 and 2 are unknown for the mask, so just deterimine those strides
      if (rank > 0) { D2 = y.dim_size(--rank); M2 = m[0].dim_size(rank) == 1 ? 0 : D3;    }
      if (rank > 0) { D1 = y.dim_size(--rank); M1 = m[0].dim_size(rank) == 1 ? 0 : D3*D2; }

      m_ptr = m[0].flat<float>().data();
    }
    while (rank > 0) { D0 *= y.dim_size(--rank); }

    OP_REQUIRES(ctx, D2 < 65536, errors::Internal("D2 < 65536: ", D2));
    OP_REQUIRES(ctx, D1 < 65536, errors::Internal("D1 < 65536: ", D1));

    Tensor* dx = NULL;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, dy.shape(), &dx));

          V* dx_ptr = (      V*)dx->flat<T>().data();
    const V* dy_ptr = (const V*)dy.flat<T>().data();
    const V*  y_ptr = (const V*)y.flat<T>().data();

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

    Benchmark* bench = nullptr;
    if (bench_)
    {
      char bench_string[256];
      sprintf(bench_string, "MaskedSoftmaxGrad (%6d,%4d,%4d,%4d) %d, %d", D0, D1, D2, D3, (uint)m.size(), (uint)sizeof(V));
      bench = new Benchmark(stream, bench_string, dy.NumElements()*3*sizeof(V), 0, bench_);
    }

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      MaskedSoftmaxGrad<V>(stream, dx_ptr, dy_ptr, y_ptr, m_ptr, D0, D1, D2, D3, M1, M2, s.scalar<float>()());

    if (bench) delete bench;
  }