Skip to content

Cast Op #1111

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

Merged
lockshaw merged 14 commits into from
Oct 12, 2023
Merged

Cast Op #1111

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
951222e
cast
KateUnger Sep 7, 2023
9dbc79f
cast cuda
KateUnger Sep 12, 2023
922e40b
cast cc
KateUnger Sep 12, 2023
a1bb222
format
KateUnger Sep 12, 2023
edb8903
Merge branch 'repo-refactor' of https://github.com/flexflow/FlexFlow …
KateUnger Oct 4, 2023
b9307b3
cast
KateUnger Oct 4, 2023
cbc89b1
cast
KateUnger Oct 4, 2023
b88dee6
add back measure_operator
KateUnger Oct 4, 2023
5c5ef0f
Merge branch 'repo-refactor' into cast_new
lockshaw Oct 6, 2023
86a5205
Fix hip and measure op cost
reyna-abhyankar Oct 7, 2023
6883962
Merge branch 'repo-refactor' into cast_new
reyna-abhyankar Oct 7, 2023
0ed0ed3
Merge branch 'repo-refactor' into cast_new
reyna-abhyankar Oct 9, 2023
4739891
Remove handle
reyna-abhyankar Oct 9, 2023
55bbd79
Merge branch 'repo-refactor' into cast_new
lockshaw Oct 12, 2023
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
20 changes: 8 additions & 12 deletions lib/kernels/include/kernels/cast_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,28 +3,24 @@

#include "kernels/accessor.h"
#include "kernels/device.h"
#include "op-attrs/ffconst.h"
#include "kernels/ff_handle.h"
#include "op-attrs/activation.h"

namespace FlexFlow {

class CastPerDeviceState : public PerDeviceOpState {
public:
CastPerDeviceState(FFHandler handle);
DataType input_data_type, output_data_type;
};

namespace Kernels {
namespace Cast {

void forward_kernel(ffStream_t stream,
CastPerDeviceState const *,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output);
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type);

void backward_kernel(ffStream_t stream,
CastPerDeviceState const *,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output);
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type);

} // namespace Cast
} // namespace Kernels
Expand Down
23 changes: 10 additions & 13 deletions lib/kernels/src/cuda/cast_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,12 @@
* limitations under the License.
*/

#include "device.h"
#include "kernels/cast_kernels.h"
#include "kernels/cuda_helper.h"
#include "kernels/datatype_dispatch.h"
#include "kernels/device.h"

namespace FlexFlow {

CastPerDeviceState::CastPerDeviceState(FFHandler handle)
: PerDeviceOpState(handle) {}

namespace Kernels {
namespace Cast {

Expand All @@ -43,7 +40,6 @@ __global__ void
template <DataType IDT, DataType ODT>
struct ForwardKernel {
void operator()(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
size_t volume = input.shape.get_volume();
Expand All @@ -55,7 +51,6 @@ struct ForwardKernel {
template <DataType IDT, DataType ODT>
struct BackwardKernel {
void operator()(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
size_t volume = input.shape.get_volume();
Expand All @@ -65,19 +60,21 @@ struct BackwardKernel {
};

void forward_kernel(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type) {
DataTypeDispatch2<ForwardKernel>{}(
m->input_data_type, m->output_data_type, stream, m, input, output);
input_type, output_type, stream, handle, input, output);
}

void backward_kernel(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type) {
DataTypeDispatch2<BackwardKernel>{}(
m->input_data_type, m->output_data_type, stream, m, input, output);
input_type, output_type, stream, handle, input, output);
}

} // namespace Cast
Expand Down
20 changes: 8 additions & 12 deletions lib/kernels/src/hip/cast_kernels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,6 @@
#include <hip/hip_runtime.h>

namespace FlexFlow {

CastPerDeviceState::CastPerDeviceState(FFHandler handle)
: PerDeviceOpState(handle) {}

namespace Kernels {
namespace Cast {

Expand All @@ -44,7 +40,6 @@ __global__ void
template <DataType IDT, DataType ODT>
struct ForwardKernel {
void operator()(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
size_t volume = input.shape.get_volume();
Expand All @@ -62,7 +57,6 @@ struct ForwardKernel {
template <DataType IDT, DataType ODT>
struct BackwardKernel {
void operator()(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
size_t volume = input.shape.get_volume();
Expand All @@ -79,19 +73,21 @@ struct BackwardKernel {
};

void forward_kernel(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type) {
DataTypeDispatch2<ForwardKernel>{}(
m->input_data_type, m->output_data_type, stream, m, input, output);
input_type, output_type, stream, input, output);
}

void backward_kernel(ffStream_t stream,
CastPerDeviceState const *m,
GenericTensorAccessorR const &input,
GenericTensorAccessorW const &output) {
GenericTensorAccessorW const &output,
DataType input_type,
DataType output_type) {
DataTypeDispatch2<BackwardKernel>{}(
m->input_data_type, m->output_data_type, stream, m, input, output);
input_type, output_type, stream, input, output);
}

} // namespace Cast
Expand Down
Loading