void Compute()

in src/transformer_op.cc [232:281]
34 lines of code
11 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override {

    const Tensor& x = ctx->input(0);
    const Tensor& s = ctx->input(1);
    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 = x.dims();
    uint D3 = x.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 = x.dim_size(--rank); M2 = m[0].dim_size(rank) == 1 ? 0 : D3;    }
      if (rank > 0) { D1 = x.dim_size(--rank); M1 = m[0].dim_size(rank) == 1 ? 0 : D3*D2; }

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

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

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

          V* y_ptr = (      V*)y->flat<T>().data();
    const V* x_ptr = (const V*)x.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, "MaskedSoftmax     (%6d,%4d,%4d,%4d) %d, %d", D0, D1, D2, D3, (uint)m.size(), (uint)sizeof(V));
      bench = new Benchmark(stream, bench_string, x.NumElements()*2*sizeof(V), 0, bench_);
    }

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

    if (bench) delete bench;
  }