diff --git a/lib/kernels/include/kernels/softmax_kernels.h b/lib/kernels/include/kernels/softmax_kernels.h index 23f0ff879d..f97f161415 100644 --- a/lib/kernels/include/kernels/softmax_kernels.h +++ b/lib/kernels/include/kernels/softmax_kernels.h @@ -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 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 + 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 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 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 &per_device_state); +void softmax_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state); -} // namespace FlexFlow::Kernels::Softmax +} // namespace FlexFlow #endif diff --git a/lib/kernels/include/kernels/softmax_kernels_cpu.h b/lib/kernels/include/kernels/softmax_kernels_cpu.h index 536a28e62c..f42a2023a7 100644 --- a/lib/kernels/include/kernels/softmax_kernels_cpu.h +++ b/lib/kernels/include/kernels/softmax_kernels_cpu.h @@ -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 +#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 diff --git a/lib/kernels/include/kernels/softmax_kernels_gpu.h b/lib/kernels/include/kernels/softmax_kernels_gpu.h index 16e98857f4..fb20e3c0f4 100644 --- a/lib/kernels/include/kernels/softmax_kernels_gpu.h +++ b/lib/kernels/include/kernels/softmax_kernels_gpu.h @@ -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 diff --git a/lib/kernels/include/kernels/softmax_per_device_state.dtg.toml b/lib/kernels/include/kernels/softmax_per_device_state.dtg.toml index abf144631e..d10d40fe74 100644 --- a/lib/kernels/include/kernels/softmax_per_device_state.dtg.toml +++ b/lib/kernels/include/kernels/softmax_per_device_state.dtg.toml @@ -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" diff --git a/lib/kernels/src/cuda/ops/softmax_kernels.cu b/lib/kernels/src/cuda/ops/softmax_kernels.cu index 1ecf42dfd8..16b94641ca 100644 --- a/lib/kernels/src/cuda/ops/softmax_kernels.cu +++ b/lib/kernels/src/cuda/ops/softmax_kernels.cu @@ -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 diff --git a/lib/kernels/src/kernels/softmax_kernels.cc b/lib/kernels/src/kernels/softmax_kernels.cc index 3cc655dc7c..f4a2d6bfcd 100644 --- a/lib/kernels/src/kernels/softmax_kernels.cc +++ b/lib/kernels/src/kernels/softmax_kernels.cc @@ -1,79 +1,93 @@ #include "kernels/softmax_kernels.h" #include "kernels/softmax_kernels_cpu.h" #include "kernels/softmax_kernels_gpu.h" -#include +#include "utils/optional.h" -namespace FlexFlow::Kernels::Softmax { +namespace FlexFlow { -std::optional 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 + softmax_init_kernel(DeviceType device_type, + SoftmaxAttrs const &attrs, + TensorShape const &input_shape, + TensorShape const &output_shape) { if (device_type == DeviceType::GPU) { - return gpu_init_kernel( - /*handle=*/handle.require_for_gpu(), - /*dim=*/dim, - /*input_n=*/input_n, - /*input_c=*/input_c, - /*input_h=*/input_h, - /*input_w=*/input_w); + return softmax_gpu_init_kernel( + /*attrs=*/attrs, + /*input_shape=*/input_shape, + /*output_shape=*/output_shape); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(handle.is_for_cpu()); return std::nullopt; } } -void forward_kernel( +void softmax_forward_kernel( device_stream_t const &stream, + device_handle_t const &handle, std::optional const &per_device_state, - float const *input_ptr, - float *output_ptr) { + SoftmaxAttrs const &attrs, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &output) { if (stream.is_gpu()) { - gpu_forward_kernel( + softmax_gpu_forward_kernel( /*stream=*/stream.require_gpu(), - /*per_device_state=*/per_device_state.value(), - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); } else { ASSERT(stream.is_cpu()); - ASSERT(per_device_state == std::nullopt); - cpu_forward_kernel( - /*input_ptr=*/input_ptr, - /*output_ptr=*/output_ptr); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + softmax_cpu_forward_kernel( + /*attrs=*/attrs, + /*input=*/input, + /*output=*/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 const &per_device_state, + SoftmaxAttrs const &attrs, + GenericTensorAccessorR const &output, + GenericTensorAccessorR const &output_grad, + GenericTensorAccessorR const &input, + GenericTensorAccessorW const &input_grad) { if (stream.is_gpu()) { - gpu_backward_kernel( + softmax_gpu_backward_kernel( /*stream=*/stream.require_gpu(), - /*output_grad_ptr=*/output_grad_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*num_elements=*/num_elements); + /*handle=*/handle.require_for_gpu(), + /*per_device_state=*/assert_unwrap(per_device_state), + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } else { ASSERT(stream.is_cpu()); - cpu_backward_kernel( - /*output_grad_ptr=*/output_grad_ptr, - /*input_grad_ptr=*/input_grad_ptr, - /*num_elements=*/num_elements); + ASSERT(handle.is_for_cpu()); + ASSERT(!per_device_state.has_value()); + softmax_cpu_backward_kernel( + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); } } -void cleanup_kernel(DeviceType device_type, - std::optional &per_device_state) { +void softmax_cleanup_kernel( + DeviceType device_type, + std::optional &per_device_state) { if (device_type == DeviceType::GPU) { - gpu_cleanup_kernel(per_device_state.value()); + softmax_gpu_cleanup_kernel(per_device_state.value()); } else { ASSERT(device_type == DeviceType::CPU); - ASSERT(per_device_state == std::nullopt); + ASSERT(!per_device_state.has_value()); } } -} // namespace FlexFlow::Kernels::Softmax +} // namespace FlexFlow diff --git a/lib/kernels/src/kernels/softmax_kernels_cpu.cc b/lib/kernels/src/kernels/softmax_kernels_cpu.cc index 20f9b68299..35f7cf9e0b 100644 --- a/lib/kernels/src/kernels/softmax_kernels_cpu.cc +++ b/lib/kernels/src/kernels/softmax_kernels_cpu.cc @@ -1,16 +1,20 @@ #include "kernels/softmax_kernels_cpu.h" #include "utils/exception.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) { NOT_IMPLEMENTED(); } -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) { NOT_IMPLEMENTED(); } -} // namespace FlexFlow::Kernels::Softmax +} // namespace FlexFlow diff --git a/lib/kernels/test/src/kernels/softmax_kernel_gpu.cc b/lib/kernels/test/src/kernels/softmax_kernel_gpu.cc new file mode 100644 index 0000000000..3588f17d93 --- /dev/null +++ b/lib/kernels/test/src/kernels/softmax_kernel_gpu.cc @@ -0,0 +1,203 @@ +#include "internal/test_utils.h" +#include "kernels/create_accessor_with_contents.h" +#include "kernels/format_accessor_contents.h" +#include "kernels/softmax_kernels_gpu.h" +#include "test/utils/doctest/check_kv.h" +#include + +using namespace ::FlexFlow; + +TEST_SUITE(FF_CUDA_TEST_SUITE) { + TEST_CASE("Softmax Kernel (GPU)") { + ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( + /*workSpaceSize=*/1024 * 1024, + /*allowTensorOpMathConversion=*/true); + ManagedFFStream managed_stream{}; + + Allocator allocator = create_local_cuda_memory_allocator(); + + GenericTensorAccessorR input = create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + allocator); + + SUBCASE("softmax_gpu_forward_kernel(dim = 0)") { + SoftmaxAttrs attrs{ + /*dim=*/ff_dim_t{0_n}, + }; + + // Intentionally randomize this tensor so we can be confident we never read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(input.shape, allocator); + + SoftmaxPerDeviceState per_device_state = + softmax_gpu_init_kernel(attrs, input.shape, output.shape); + + softmax_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {0.0023548827, 0.0023555316, 0.0023556193}, + {0.04729908, 0.047312114, 0.047313876}, + {0.9500274, 0.95028925, 0.95032465}, + {0.00031869867, 4.3143067e-05, 5.8389965e-06}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + SUBCASE("softmax_gpu_forward_kernel(dim = 1)") { + SoftmaxAttrs attrs{ + /*dim=*/ff_dim_t{1_n}, + }; + + // Intentionally randomize this tensor so we can be confident we never read it + GenericTensorAccessorW output = + create_random_filled_accessor_w(input.shape, allocator); + + SoftmaxPerDeviceState per_device_state = + softmax_gpu_init_kernel(attrs, input.shape, output.shape); + + softmax_gpu_forward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*input=*/input, + /*output=*/output); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {0.090030566, 0.24472848, 0.6652409}, + {0.090030566, 0.24472848, 0.6652409}, + {0.090030566, 0.24472848, 0.6652409}, + {0.6652409, 0.24472848, 0.090030566}, + }, + allocator); + + CHECK_MESSAGE(accessors_are_equal(output, correct), + check_kv("output", format_accessor_w_contents(output))); + } + + SUBCASE("softmax_gpu_backward_kernel(dim = 0)") { + SoftmaxAttrs attrs{ + /*dim=*/ff_dim_t{0_n}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {0.0023548827, 0.0023555316, 0.0023556193}, + {0.04729908, 0.047312114, 0.047313876}, + {0.9500274, 0.95028925, 0.95032465}, + {0.00031869867, 4.3143067e-05, 5.8389965e-06}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + allocator); + GenericTensorAccessorW input_grad = + create_zero_filled_accessor_w(input.shape, allocator); + + SoftmaxPerDeviceState per_device_state = + softmax_gpu_init_kernel(attrs, input.shape, output.shape); + + softmax_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-0.013755869, -0.013764547, -0.013765898}, + {-0.13439676, -0.13453196, -0.1345538}, + {0.1506511, 0.14872104, 0.14838853}, + {-0.0024990516, -0.00042467876, -6.915622e-05}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + SUBCASE("softmax_gpu_backward_kernel(dim = 1)") { + SoftmaxAttrs attrs{ + /*dim=*/ff_dim_t{1_n}, + }; + + GenericTensorAccessorR output = create_2d_accessor_r_with_contents( + { + {0.090030566, 0.24472848, 0.6652409}, + {0.090030566, 0.24472848, 0.6652409}, + {0.090030566, 0.24472848, 0.6652409}, + {0.6652409, 0.24472848, 0.090030566}, + }, + allocator); + + GenericTensorAccessorR output_grad = + create_2d_accessor_r_with_contents( + { + {1, 2, 3}, + {4, 5, 6}, + {7, 8, 9}, + {-1, -2, -3}, + }, + allocator); + GenericTensorAccessorW input_grad = + create_zero_filled_accessor_w(input.shape, allocator); + + SoftmaxPerDeviceState per_device_state = + softmax_gpu_init_kernel(attrs, input.shape, output.shape); + + softmax_gpu_backward_kernel( + /*stream=*/managed_stream.raw_stream(), + /*handle=*/managed_handle.raw_handle(), + /*per_device_state=*/per_device_state, + /*attrs=*/attrs, + /*output=*/output, + /*output_grad=*/output_grad, + /*input=*/input, + /*input_grad=*/input_grad); + + GenericTensorAccessorR correct = + create_2d_accessor_r_with_contents( + { + {-0.14181706, -0.14077029, 0.28258762}, + {-0.14181706, -0.14077029, 0.28258762}, + {-0.14181702, -0.14077017, 0.28258792}, + {0.28258738, -0.14077038, -0.1418171}, + }, + allocator); + + CHECK_MESSAGE( + accessors_are_equal(input_grad, correct), + check_kv("input_grad", format_accessor_w_contents(input_grad))); + } + } +} diff --git a/lib/kernels/test/src/test_softmax_kernel.cc b/lib/kernels/test/src/test_softmax_kernel.cc deleted file mode 100644 index ca94bf58d1..0000000000 --- a/lib/kernels/test/src/test_softmax_kernel.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "internal/test_utils.h" -#include "kernels/softmax_kernels_gpu.h" -#include - -using namespace ::FlexFlow; - -TEST_SUITE(FF_CUDA_TEST_SUITE) { - TEST_CASE("Test Softmax Kernel Operations") { - nonnegative_int input_n = 1_n; - nonnegative_int input_c = 1_n; - nonnegative_int input_h = 1_n; - nonnegative_int input_w = 100_n; - nonnegative_int channels = 100_n; - - ManagedPerDeviceFFHandle managed_handle = initialize_single_gpu_handle( - /*workSpaceSize=*/1024 * 1024, - /*allowTensorOpMathConversion=*/true); - ManagedFFStream managed_stream{}; - - Allocator allocator = create_local_cuda_memory_allocator(); - - TensorShape input_shape = TensorShape{ - TensorDims{FFOrdered{100_p}}, - DataType::FLOAT, - }; - TensorShape output_shape = input_shape; - - SoftmaxPerDeviceState state = - Kernels::Softmax::gpu_init_kernel(managed_handle.raw_handle(), - ff_dim_t{3_n}, - input_n.unwrap_nonnegative(), - channels.unwrap_nonnegative(), - input_h.unwrap_nonnegative(), - input_w.unwrap_nonnegative()); - - GenericTensorAccessorW output_accessor = - create_random_filled_accessor_w(output_shape, allocator); - - SUBCASE("gpu_forward_kernel") { - GenericTensorAccessorW input_accessor = - create_random_filled_accessor_w(input_shape, allocator); - - Kernels::Softmax::gpu_forward_kernel(managed_stream.raw_stream(), - state, - input_accessor.get_float_ptr(), - output_accessor.get_float_ptr()); - - CHECK(contains_non_zero(output_accessor)); - } - - SUBCASE("gpu_backward_kernel") { - GenericTensorAccessorR output_grad_accessor = - create_random_filled_accessor_r(output_shape, allocator); - GenericTensorAccessorW input_grad_accessor = - allocator.allocate_tensor(input_shape); - - Kernels::Softmax::gpu_backward_kernel( - managed_stream.raw_stream(), - output_grad_accessor.get_float_ptr(), - input_grad_accessor.get_float_ptr(), - get_num_elements(output_grad_accessor.shape.dims) - .int_from_positive_int()); - - CHECK(contains_non_zero(input_grad_accessor)); - } - } -} diff --git a/lib/task-spec/src/task-spec/ops/impl/softmax.cc b/lib/task-spec/src/task-spec/ops/impl/softmax.cc index 66693913e6..b777f60fb2 100644 --- a/lib/task-spec/src/task-spec/ops/impl/softmax.cc +++ b/lib/task-spec/src/task-spec/ops/impl/softmax.cc @@ -22,29 +22,16 @@ namespace FlexFlow { -using namespace FlexFlow::Kernels::Softmax; - static DeviceSpecificPerDeviceOpState init_task_impl(TaskArgumentAccessor const &acc) { - device_handle_t handle = acc.get_ff_handle(); DeviceType kernel_device_type = acc.get_kernel_device_type(); SoftmaxAttrs attrs = acc.get_op_attrs().require_softmax(); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - positive_int output_w = dim_at_idx(output.shape.dims, legion_dim_t{0_n}); - positive_int output_h = dim_at_idx(output.shape.dims, legion_dim_t{1_n}); - positive_int output_c = dim_at_idx(output.shape.dims, legion_dim_t{2_n}); - positive_int output_n = dim_at_idx(output.shape.dims, legion_dim_t{3_n}); + TensorShape input_shape = acc.get_tensor_shape(TensorSlotName::INPUT); + TensorShape output_shape = acc.get_tensor_shape(TensorSlotName::OUTPUT); std::optional per_device_state = - init_kernel(kernel_device_type, - handle, - attrs.dim, - output_n.int_from_positive_int(), - output_c.int_from_positive_int(), - output_h.int_from_positive_int(), - output_w.int_from_positive_int()); + softmax_init_kernel(kernel_device_type, attrs, input_shape, output_shape); return DeviceSpecificPerDeviceOpState{ acc.make_device_specific(per_device_state), @@ -53,48 +40,60 @@ static DeviceSpecificPerDeviceOpState static std::optional forward_task_impl(TaskArgumentAccessor const &acc) { + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW output = + acc.get_tensor(TensorSlotName::OUTPUT); + SoftmaxAttrs attrs = acc.get_op_attrs().require_softmax(); + + device_handle_t handle = acc.get_ff_handle(); + ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - SoftmaxPerDeviceState per_device_state = - acc.get_per_device_op_state().require_softmax().value(); - - auto input = acc.get_tensor(TensorSlotName::INPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); + std::optional per_device_state = + acc.get_per_device_op_state().require_softmax(); - return profile(forward_kernel, + return profile(softmax_forward_kernel, profiling, kernel_device_type, "[Softmax] forward_time = {:.2lf}ms\n", + handle, per_device_state, - input.get_float_ptr(), - output.get_float_ptr()); + attrs, + input, + output); } static std::optional backward_task_impl(TaskArgumentAccessor const &acc) { + GenericTensorAccessorR input = + acc.get_tensor(TensorSlotName::INPUT); + GenericTensorAccessorW input_grad = + acc.get_tensor_grad(TensorSlotName::INPUT); + GenericTensorAccessorR output = + acc.get_tensor(TensorSlotName::OUTPUT); + GenericTensorAccessorR output_grad = + acc.get_tensor_grad(TensorSlotName::OUTPUT); + + SoftmaxAttrs attrs = acc.get_op_attrs().require_softmax(); + device_handle_t handle = acc.get_ff_handle(); + ProfilingSettings profiling = acc.get_profiling_settings(); DeviceType kernel_device_type = acc.get_kernel_device_type(); - SoftmaxPerDeviceState per_device_state = - acc.get_per_device_op_state().require_softmax().value(); - - auto input_grad = acc.get_tensor_grad(TensorSlotName::INPUT); - auto input = acc.get_tensor(TensorSlotName::INPUT); - assert(input_grad.shape == input.shape); + std::optional per_device_state = + acc.get_per_device_op_state().require_softmax(); - auto output_grad = - acc.get_tensor_grad(TensorSlotName::OUTPUT); - auto output = acc.get_tensor(TensorSlotName::OUTPUT); - - assert(output_grad.shape == output.shape); - - return profile( - backward_kernel, - profiling, - kernel_device_type, - "[Softmax] backward_time = {:.2lf}ms\n", - output_grad.get_float_ptr(), - input_grad.get_float_ptr(), - get_num_elements(output_grad.shape.dims).int_from_positive_int()); + return profile(softmax_backward_kernel, + profiling, + kernel_device_type, + "[Softmax] backward_time = {:.2lf}ms\n", + handle, + per_device_state, + attrs, + output, + output_grad, + input, + input_grad); } TaskImplFunction get_softmax_init_task_impl() { @@ -109,4 +108,4 @@ TaskImplFunction get_softmax_bwd_task_impl() { return TaskImplFunction{FwdBwdOpTaskImplFunction{backward_task_impl}}; } -}; // namespace FlexFlow +} // namespace FlexFlow