void Compute()

in src/embedding_op.cc [61:106]

• 37 lines of code
• 7 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override
  {
    if (SMs_ == 0)
      SMs_ = GetCountSMs();

    const Tensor& emb = ctx->input(0);
    const Tensor& idx = ctx->input(1);

    OP_REQUIRES(ctx, emb.dim_size(0) == ctx->input(2).scalar<int32>()(), errors::InvalidArgument("Bad emb channels arg"));

    int C    = emb.dim_size(0);
    int K    = emb.dim_size(1);
    int rank = idx.dims();
    int nIdx = 1;
    TensorShape shape;
    for (int i = 0; i < rank; ++i)
    {
      int dim = idx.dim_size(i);
      nIdx *= dim;
      shape.AddDim(dim);
    }
    shape.AddDim(K);

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

           V*   y_ptr = (V*)y->flat<T>().data();
    const  V* emb_ptr = (const V*)emb.flat<T>().data();
    const TI* idx_ptr = idx.flat<TI>().data();

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

    Benchmark* bench = nullptr;
    if (bench_)
    {
      char bench_string[256];
      sprintf(bench_string, "EmbeddingLookup     nIdx:%7d, C:%5d, K:%4d", nIdx, C, K);
      bench = new Benchmark(stream, bench_string, nIdx*sizeof(TI) + 2*nIdx*K*sizeof(V), 0, bench_);
    }

    int repeat = bench_ ? bench_ : 1;
    for (int i = 0; i < repeat; i++)
      EmbeddingLookup<TI,V>(stream, SMs_, y_ptr, idx_ptr, emb_ptr, nIdx, C, K);

    if (bench) delete bench;
  }