Skip to content

Update and sync softmax interface to modern style #1668

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
elliottslaughter wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Update and sync softmax interface to modern style #1668

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 27 additions & 19 deletions lib/kernels/include/kernels/softmax_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,36 +1,44 @@
#ifndef _FLEXFLOW_OPS_KERNELS_SOFTMAX_KERNELS_H
#define _FLEXFLOW_OPS_KERNELS_SOFTMAX_KERNELS_H

#include "kernels/accessor.h"
#include "kernels/device_handle_t.dtg.h"
#include "kernels/device_stream_t.dtg.h"
#include "kernels/ff_handle.h"
#include "kernels/softmax_per_device_state.dtg.h"
#include "op-attrs/ops/softmax_attrs.dtg.h"
#include "op-attrs/tensor_shape.dtg.h"
#include "pcg/device_type.dtg.h"

namespace FlexFlow::Kernels::Softmax {
namespace FlexFlow {

std::optional<SoftmaxPerDeviceState> init_kernel(DeviceType device_type,
device_handle_t const &handle,
ff_dim_t dim,
int input_n,
int input_c,
int input_h,
int input_w);
std::optional<SoftmaxPerDeviceState>
softmax_init_kernel(DeviceType device_type,
SoftmaxAttrs const &attrs,
TensorShape const &input_shape,
TensorShape const &output_shape);

void forward_kernel(
void softmax_forward_kernel(
device_stream_t const &stream,
device_handle_t const &handle,
std::optional<SoftmaxPerDeviceState> const &per_device_state,
float const *input_ptr,
float *output_ptr);
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output);

void backward_kernel(device_stream_t const &stream,
float const *output_grad_ptr,
float *input_grad_ptr,
size_t num_elements);
void softmax_backward_kernel(
device_stream_t const &stream,
device_handle_t const &handle,
std::optional<SoftmaxPerDeviceState> const &per_device_state,
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &input_grad);

void cleanup_kernel(DeviceType device_type,
std::optional<SoftmaxPerDeviceState> &per_device_state);
void softmax_cleanup_kernel(
DeviceType device_type,
std::optional<SoftmaxPerDeviceState> &per_device_state);

} // namespace FlexFlow::Kernels::Softmax
} // namespace FlexFlow

#endif
19 changes: 12 additions & 7 deletions lib/kernels/include/kernels/softmax_kernels_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,21 @@
#ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SOFTMAX_KERNELS_CPU_H
#define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SOFTMAX_KERNELS_CPU_H

#include <cstddef>
#include "kernels/accessor.h"
#include "op-attrs/ops/softmax_attrs.dtg.h"

namespace FlexFlow::Kernels::Softmax {
namespace FlexFlow {

void cpu_forward_kernel(float const *input_ptr, float *output_ptr);
void softmax_cpu_forward_kernel(SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output);

void cpu_backward_kernel(float const *output_grad_ptr,
float *input_grad_ptr,
size_t num_elements);
void softmax_cpu_backward_kernel(SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &input_grad);

} // namespace FlexFlow::Kernels::Softmax
} // namespace FlexFlow

#endif
41 changes: 22 additions & 19 deletions lib/kernels/include/kernels/softmax_kernels_gpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,32 +1,35 @@
#ifndef _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SOFTMAX_KERNELS_GPU_H
#define _FLEXFLOW_LIB_KERNELS_INCLUDE_KERNELS_SOFTMAX_KERNELS_GPU_H

#include "kernels/accessor.h"
#include "kernels/device.h"
#include "kernels/ff_handle.h"
#include "kernels/softmax_per_device_state.dtg.h"
#include "op-attrs/ff_dim_t.dtg.h"
#include "op-attrs/ops/softmax_attrs.dtg.h"

namespace FlexFlow::Kernels::Softmax {
namespace FlexFlow {

SoftmaxPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle,
ff_dim_t dim,
int input_n,
int input_c,
int input_h,
int input_w);
SoftmaxPerDeviceState softmax_gpu_init_kernel(SoftmaxAttrs const &attrs,
TensorShape const &input_shape,
TensorShape const &output_shape);

void gpu_forward_kernel(ffStream_t stream,
SoftmaxPerDeviceState const &per_device_state,
float const *input_ptr,
float *output_ptr);
void softmax_gpu_forward_kernel(ffStream_t stream,
PerDeviceFFHandle const &handle,
SoftmaxPerDeviceState const &per_device_state,
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output);

void gpu_backward_kernel(ffStream_t stream,
float const *output_grad_ptr,
float *input_grad_ptr,
size_t num_elements);
void softmax_gpu_backward_kernel(ffStream_t stream,
PerDeviceFFHandle const &handle,
SoftmaxPerDeviceState const &per_device_state,
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &input_grad);

void gpu_cleanup_kernel(SoftmaxPerDeviceState &per_device_state);
void softmax_gpu_cleanup_kernel(SoftmaxPerDeviceState &per_device_state);

} // namespace FlexFlow::Kernels::Softmax
} // namespace FlexFlow

#endif
10 changes: 5 additions & 5 deletions lib/kernels/include/kernels/softmax_per_device_state.dtg.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,13 @@ includes = [
]

[[fields]]
name = "handle"
type = "::FlexFlow::PerDeviceFFHandle"
name = "inputTensor"
type = "ffTensorDescriptor_t"

[[fields]]
name = "inputTensor"
name = "outputTensor"
type = "ffTensorDescriptor_t"

[[fields]]
name = "dim"
type = "::FlexFlow::ff_dim_t"
name = "outputGradTensor"
type = "ffTensorDescriptor_t"
123 changes: 81 additions & 42 deletions lib/kernels/src/cuda/ops/softmax_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,72 +15,111 @@

#include "internal/device.h"
#include "kernels/softmax_kernels_gpu.h"
#include "op-attrs/ff_dim_t.h"
#include "op-attrs/tensor_dims.h"
#include "op-attrs/tensor_shape.h"
#include "utils/exception.h"

namespace FlexFlow {

namespace Kernels {
namespace Softmax {

SoftmaxPerDeviceState gpu_init_kernel(PerDeviceFFHandle const &handle,
ff_dim_t dim,
int input_n,
int input_c,
int input_h,
int input_w) {
SoftmaxPerDeviceState softmax_gpu_init_kernel(SoftmaxAttrs const &attrs,
TensorShape const &input_shape,
TensorShape const &output_shape) {
ffTensorDescriptor_t inputTensor;
ffTensorDescriptor_t outputTensor;
ffTensorDescriptor_t outputGradTensor;

TensorShape shape = require_same(input_shape, output_shape);

positive_int num_outer_elements =
get_num_elements(slice_tensor_dims(shape.dims, ff_dim_t{0_n}, attrs.dim));
positive_int softmax_dim_size = dim_at_idx(shape.dims, attrs.dim);
positive_int num_inner_elements = get_num_elements(
slice_tensor_dims(shape.dims, add_to_ff_dim(attrs.dim, 1), std::nullopt));

checkCUDNN(cudnnCreateTensorDescriptor(&inputTensor));
checkCUDNN(cudnnSetTensor4dDescriptor(inputTensor,
CUDNN_TENSOR_NCHW,
CUDNN_DATA_FLOAT,
input_n,
input_c,
input_h,
input_w));
checkCUDNN(cudnnSetTensor4dDescriptor(
inputTensor,
CUDNN_TENSOR_NCHW,
ff_to_cudnn_datatype(shape.data_type),
/*n=*/num_outer_elements.int_from_positive_int(),
/*c=*/softmax_dim_size.int_from_positive_int(),
/*h=*/num_inner_elements.int_from_positive_int(),
/*w=*/1));

SoftmaxPerDeviceState per_device_state = SoftmaxPerDeviceState{
/*handle=*/handle,
checkCUDNN(cudnnCreateTensorDescriptor(&outputTensor));
checkCUDNN(cudnnSetTensor4dDescriptor(
outputTensor,
CUDNN_TENSOR_NCHW,
ff_to_cudnn_datatype(shape.data_type),
/*n=*/num_outer_elements.int_from_positive_int(),
/*c=*/softmax_dim_size.int_from_positive_int(),
/*h=*/num_inner_elements.int_from_positive_int(),
/*w=*/1));

checkCUDNN(cudnnCreateTensorDescriptor(&outputGradTensor));
checkCUDNN(cudnnSetTensor4dDescriptor(
outputGradTensor,
CUDNN_TENSOR_NCHW,
ff_to_cudnn_datatype(shape.data_type),
/*n=*/num_outer_elements.int_from_positive_int(),
/*c=*/softmax_dim_size.int_from_positive_int(),
/*h=*/num_inner_elements.int_from_positive_int(),
/*w=*/1));

return SoftmaxPerDeviceState{
/*inputTensor=*/inputTensor,
/*dim=*/dim,
/*outputTensor=*/outputTensor,
/*outputGradTensor=*/outputGradTensor,
};
return per_device_state;
}

void gpu_forward_kernel(cudaStream_t stream,
SoftmaxPerDeviceState const &m,
float const *input_ptr,
float *output_ptr) {
checkCUDNN(cudnnSetStream(m.handle.dnn, stream));
void softmax_gpu_forward_kernel(ffStream_t stream,
PerDeviceFFHandle const &handle,
SoftmaxPerDeviceState const &per_device_state,
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
checkCUDNN(cudnnSetStream(handle.dnn, stream));

float alpha = 1.0f, beta = 0.0f;
checkCUDNN(cudnnSoftmaxForward(m.handle.dnn,
checkCUDNN(cudnnSoftmaxForward(handle.dnn,
CUDNN_SOFTMAX_ACCURATE,
CUDNN_SOFTMAX_MODE_CHANNEL,
&alpha,
m.inputTensor,
input_ptr,
per_device_state.inputTensor,
input.get_float_ptr(),
&beta,
m.inputTensor,
output_ptr));
per_device_state.inputTensor,
output.get_float_ptr()));
}

void gpu_backward_kernel(cudaStream_t stream,
float const *output_grad_ptr,
float *input_grad_ptr,
size_t num_elements) {
void softmax_gpu_backward_kernel(ffStream_t stream,
PerDeviceFFHandle const &handle,
SoftmaxPerDeviceState const &per_device_state,
SoftmaxAttrs const &attrs,
GenericTensorAccessorR const &output,
GenericTensorAccessorR const &output_grad,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &input_grad) {
checkCUDNN(cudnnSetStream(handle.dnn, stream));

checkCUDA(cudaMemcpyAsync(input_grad_ptr,
output_grad_ptr,
num_elements * sizeof(float),
cudaMemcpyDeviceToDevice,
stream));
float alpha = 1.0f, beta = 0.0f;
checkCUDNN(cudnnSoftmaxBackward(handle.dnn,
CUDNN_SOFTMAX_ACCURATE,
CUDNN_SOFTMAX_MODE_CHANNEL,
&alpha,
per_device_state.inputTensor,
output.get_float_ptr(),
per_device_state.inputTensor,
output_grad.get_float_ptr(),
&beta,
per_device_state.inputTensor,
input_grad.get_float_ptr()));
}

void gpu_cleanup_kernel(SoftmaxPerDeviceState &) {
void softmax_gpu_cleanup_kernel(SoftmaxPerDeviceState &per_device_state) {
NOT_IMPLEMENTED();
}

} // namespace Softmax
} // namespace Kernels
} // namespace FlexFlow
Loading
Loading