void Compute()

in src/ew_op.cc [615:682]
52 lines of code
9 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override
  {
    const Tensor& x = ctx->input(0);
    const Tensor& m = ctx->input(1);
    float keep_prob = ctx->input(2).scalar<float>()();
    float scale     = 1.0f / keep_prob;

    // process all the shape logic on the first call
    if (SMs == 0)
    {
      SMs = GetCountSMs();

      size = x.shape().num_elements();

      // Just treat the mask as simple 1d flat shape
      if (mask_shape.empty())
      {
        OP_REQUIRES(ctx, m.shape().num_elements() == CEIL_DIV(size, 32), errors::Internal("ApplyDropoutMaskOp: bad mask shape (size)"));
        rank = 1;
        ms.stride[0] = 1;
      }
      // Setup mask broadcasting support
      else
      {
        rank = x.dims();
        OP_REQUIRES(ctx, rank == mask_shape.size(), errors::Internal("ApplyDropoutMaskOp: bad mask shape (rank)"));
        OP_REQUIRES(ctx, rank >= 1 && rank <= 5,    errors::Internal("ApplyDropoutMaskOp: only rank 1-5 tensors currently supported: ", rank));

        // Check mask size
        int mask_size = 1;
        for (int i = 0; i < rank; i++)
          mask_size *= mask_shape[i];
        OP_REQUIRES(ctx, m.shape().num_elements() == CEIL_DIV(mask_size, 32), errors::Internal("ApplyDropoutMaskOp: bad mask shape (size)"));

        // Build the strides
        int dim = rank - 1;
        xs.stride[dim] = ms.stride[dim] = 1;
        while (--dim >= 0)
        {
          ms.stride[dim] = mask_shape[dim+1] * ms.stride[dim+1];
          xs.stride[dim] = x.dim_size(dim+1) * xs.stride[dim+1];
        }

        // Update the mask strides to support broadcast
        for (int i = 0; i < rank; i++)
        {
          if (mask_shape[i] != x.dim_size(i))
          {
            if (mask_shape[i] == 1)
              ms.stride[i] = 0;
            else
              OP_REQUIRES(ctx, false, errors::Internal("ApplyDropoutMaskOp: bad mask shape (dims)"));
          }
        }
      }
    }

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

            V1* y_ptr = (        V1*)y->flat<T>().data();
    const   V1* x_ptr = (const   V1*)x.flat<T>().data();
    const uint* m_ptr = (const uint*)m.flat<int32>().data();

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

    ApplyDropoutMask<V1,V4,V8>(stream, SMs, y_ptr, x_ptr, m_ptr, scale, size, rank, xs, ms);
  }