Tensorflow Architecture: Parallism

#ifdef INTEL_MKL
template <typename Device, typename T>
class MklSoftmaxOp : public OpKernel {
   public:
    ~MklSoftmaxOp() {}

    explicit MklSoftmaxOp(OpKernelConstruction* context) : OpKernel(context) {}

    void Compute(OpKernelContext* context) override try {
        // ...
        // NOTE: softmax_fwd 是 mkldnn::softmax_forward
        softmax_fwd->Execute(src_data, dst_data);
    }
};
#endif

#if defined(GOOGLE_CUDA)
// NOTE: EIGEN_USE_GPU 会
#define EIGEN_USE_GPU
REGISTER_KERNEL_BUILDER(
    Name("ArgMax")
        .Device(DEVICE_GPU)
        .TypeConstraint<type>("T")
        .HostMemory("dimension"),
    ArgMaxOp<GPUDevice, type, int64>);
#endif

// NOTE: ArgMaxOp 是 CPU/GPU 通用的, 不需要包在 GOOGLE_CUDA 宏中
template <typename Device, typename T, typename Tout>
class ArgMaxOp
    : public ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> > {
   public:
    explicit ArgMaxOp(OpKernelConstruction* context)
        : ArgOp<Device, T, Tout, functor::ArgMax<Device, T, Tout> >(context) {}
};

template <typename Device, typename T, typename Tout, typename ArgFunctor>
class ArgOp : public OpKernel {
   public:
    explicit ArgOp(OpKernelConstruction* context) : OpKernel(context) {}

    void Compute(OpKernelContext* context) override {
        // ...

#define HANDLE_DIM(NDIM) \
    case NDIM:           \
        // NOTE: 最终会调用 ArgFunctor::Reduce1 等, 例如 ArgMax::Reduce1 \
        ArgFunctor::Reduce##NDIM(                                           \
            context->eigen_device<Device>(), input.tensor<T, NDIM>(), axis, \
            output->tensor<Tout, NDIM - 1>());                              \
        break;

        switch (input_dims) {
            HANDLE_DIM(1);
            HANDLE_DIM(2);
            // ...
        }
    }
#undef HANDLE_DIM
};

template <typename Device, typename T, typename Tout>
struct ArgMax {
#define DECLARE_COMPUTE_SPEC(Dims)                                    \
    EIGEN_ALWAYS_INLINE static void Reduce##Dims(                     \
        const Device& d, typename TTypes<T, Dims>::ConstTensor input, \
        const int32 dimension,                                        \
        typename TTypes<Tout, Dims - 1>::Tensor output) {             \
        output.device(d) = input.argmax(dimension).template cast<Tout>();
}

DECLARE_COMPUTE_SPEC(1);
DECLARE_COMPUTE_SPEC(2);
#undef DECLARE_COMPUTE_SPEC
}
;

// NOTE: input 是 Eigen::TensorMap, 这里会调用 engien 的 argmax 方法,
// 最终会调用到 TensorReductionGpu.h 中相应的 kernel 例如 FullReductionKernel

#if GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(
    Name("MatMul").Device(DEVICE_GPU).TypeConstraint<T>("T"),
    MatMulOp<GPUDevice, T, true /* cublas, true by default */>);
#endif
// --------------------
template <typename Device, typename T, bool USE_CUBLAS>
class MatMulOp : public OpKernel {
   public:
    void Compute(OpKernelContext* ctx) override {
        // ...
        LaunchMatMul<Device, float, USE_CUBLAS>::launch(
            ctx, a_float, b_float, dim_pair, &algorithms_, use_autotune_,
            &out_float);
        // ...
    }
};
// --------------------
#if GOOGLE_CUDA
template <typename T>
struct LaunchMatMul<GPUDevice, T, true /* USE_CUBLAS */> {
    static void launch() {
        // ...
        auto* stream = ctx->op_device_context()->stream();
        // ...
        bool blas_launch_status =
            stream
                ->ThenBlasGemm(
                    blas_transpose_b, blas_transpose_a, n, m, k, 1.0f, b_ptr,
                    transpose_b ? k : n, a_ptr, transpose_a ? m : k, 0.0f,
                    &c_ptr, n)
                .ok();
        // ...
    }
    // NOTE: ThenBlasGemm 是 stream_executor 的接口
#endif

template <typename Device, typename T>
class SoftmaxOp : public OpKernel {
   public:
    explicit SoftmaxOp(OpKernelConstruction* context) : OpKernel(context) {
        log_ = absl::StartsWith(type_string(), "Log");
    }

    void Compute(OpKernelContext* context) override {
        // ...
        if (logits_in.NumElements() > 0) {
            // NOTE: 这里的 functor 会最终调用到 SoftmaxEigenImpl<SyclDevice>
            functor::SoftmaxFunctor<Device, T> functor;
            functor(
                context->eigen_device<Device>(), logits_in.flat_inner_dims<T>(),
                softmax_out->flat_inner_dims<T>(), log_);
        }
    }
};
// --------------------
template <typename Device, typename T>
struct SoftmaxEigenImpl {
    static void Compute(
        const Device& d, typename TTypes<T>::ConstMatrix logits,
        typename TTypes<T>::Matrix softmax, const bool log) {
        // ...
        // logits 是 Eigen TensorMap
        auto shifted_logits =
            (logits - logits.maximum(along_class)
                          .eval()
                          .reshape(batch_by_one)
                          .broadcast(one_by_class));
        // ...
    }
};
// --------------------
// NOTE: 注册
REGISTER_KERNEL_BUILDER(
    Name("Softmax").Device(DEVICE_SYCL).TypeConstraint<float>("T"),
    SoftmaxOp<SYCLDevice, float>);

template <typename T>
struct LaunchAvgPooling3dGradOp<SYCLDevice, T> {
    static void launch(...) {
        // ...
        device.sycl_queue().submit([&](cl::sycl::handler& cgh) {
            auto input_backprop_access =
                input_backprop_buffer
                    .template get_access<cl::sycl::access::mode::read>(cgh);
            auto output_backprop_access =
                output_backprop_buffer
                    .template get_access<cl::sycl::access::mode::write>(cgh);
            AvgPool3DGradSYCL<T> functor(
                depth, batch, in_planes, in_rows, in_cols, output_shape, window,
                stride, padding, input_backprop_access, output_backprop_access);

            cgh.parallel_for(cl::sycl::range<1>(num_threads), functor);
        });
    }
};
// --------------------
#ifdef TENSORFLOW_USE_SYCL
REGISTER_KERNEL_BUILDER(
    Name("AvgPool3DGrad")
        .Device(DEVICE_SYCL)
        .TypeConstraint<T>("T")
        .HostMemory("orig_input_shape"),
    AvgPooling3dGradOp<SYCLDevice, T>);
#endif  // TENSORFLOW_USE_SYCL

class SoftmaxOp : public XlaOpKernel {
   public:
    explicit SoftmaxOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {
        log_ = absl::StartsWith(type_string(), "Log");
    }

    void Compile(XlaOpKernelContext* ctx) override {
        // ...
        xla::XlaBuilder* const b = ctx->builder();
        const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type);

        auto logits_max = xla::Reduce(
            logits, xla::MinValue(b, xla_type), max_func, {kClassDim});
        auto shifted_logits = xla::Sub(logits, logits_max, batch_dims);
        auto exp_shifted = xla::Exp(shifted_logits);
        // ...
        auto reduce = xla::Reduce(
            converted, xla::Zero(b, xla_accumulation_type),
            *ctx->GetOrCreateAdd(accumulation_type), {kClassDim});
        auto sum = XlaHelpers::ConvertElementType(reduce, type);
        // NOTE: softmax=exp/sum, sum 又来自 reduce (..., add)
        auto softmax = log_
                           ? xla::Sub(shifted_logits, xla::Log(sum), batch_dims)
                           : xla::Div(exp_shifted, sum, batch_dims);
        ctx->SetOutput(0, softmax);
    }
}

dataset(generator)
  .map(fn1,num_parallel_calls)
  .filter(fn2)
  .shuffle()
  .prefetch(buffer_size);