Skip to content

Commit a8d3cc0

Browse files
ronnie-devtechronnie
ronnie-devtech
and
ronnie
authored
Adapt MUSA 5.1 and TokenMixer perf paths (#287)
Co-authored-by: ronnie <ronnie@mt.com>
1 parent 5854cd5 commit a8d3cc0

16 files changed

Lines changed: 1856 additions & 502 deletions

‎CMakeLists.txt‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,10 +109,15 @@ endif()
109109
set(TF_COMPILE_FLAGS "${TF_COMPILE_FLAGS} -DMUSA_DISABLE_DEBUG_LOGGING -DNDEBUG -DMUSA_DISABLE_AUTO_TRACE -include ${CMAKE_CURRENT_SOURCE_DIR}/musa_ext/tf_status_compat.h")
110110

111111
# Include paths
112-
include_directories("${TF_INC}" "${MUSA_PATH}/include")
112+
#
113+
# MUSA 5.1 ships both a top-level C API header (`include/mudnn.h`) and the
114+
# C++ API headers under `include/mudnncxx/`. This plugin uses the C++ API
115+
# (`musa::dnn::*`), so put mudnncxx first to keep existing `#include <mudnn.h>`
116+
# sites resolving to the C++ aggregate header.
113117
if(EXISTS "${MUSA_PATH}/include/mudnncxx")
114118
include_directories("${MUSA_PATH}/include/mudnncxx")
115119
endif()
120+
include_directories("${TF_INC}" "${MUSA_PATH}/include")
116121

117122
# Include TensorFlow C++ core headers for common_runtime
118123
# These are required for accessing BFCAllocator and device-specific headers

‎musa_ext/kernels/array/musa_fill_kernel.mu‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
#include <musa_runtime.h>
2+
#include <stdint.h>
3+
4+
#include "tensorflow/core/framework/types.h"
5+
6+
namespace tensorflow {
7+
namespace musa {
8+
namespace {
9+
10+
constexpr int kThreadsPerBlock = 256;
11+
12+
inline int64_t CeilDiv(int64_t x, int64_t y) {
13+
return (x + y - 1) / y;
14+
}
15+
16+
template <typename T>
17+
__global__ void FillKernel(T* out, T value, int64_t n) {
18+
const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
19+
if (idx < n) {
20+
out[idx] = value;
21+
}
22+
}
23+
24+
template <typename T>
25+
void LaunchFill(T* out, T value, int64_t n, musaStream_t stream) {
26+
if (n <= 0) {
27+
return;
28+
}
29+
const int64_t blocks = CeilDiv(n, kThreadsPerBlock);
30+
FillKernel<T><<<blocks, kThreadsPerBlock, 0, stream>>>(out, value, n);
31+
}
32+
33+
} // namespace
34+
35+
extern "C" {
36+
37+
void LaunchMusaFill_float(float* out, float value, int64_t n,
38+
musaStream_t stream) {
39+
LaunchFill<float>(out, value, n, stream);
40+
}
41+
42+
void LaunchMusaFill_double(double* out, double value, int64_t n,
43+
musaStream_t stream) {
44+
LaunchFill<double>(out, value, n, stream);
45+
}
46+
47+
void LaunchMusaFill_int32(int32* out, int32 value, int64_t n,
48+
musaStream_t stream) {
49+
LaunchFill<int32>(out, value, n, stream);
50+
}
51+
52+
void LaunchMusaFill_int64(int64* out, int64 value, int64_t n,
53+
musaStream_t stream) {
54+
LaunchFill<int64>(out, value, n, stream);
55+
}
56+
57+
void LaunchMusaFill_half(Eigen::half* out, Eigen::half value, int64_t n,
58+
musaStream_t stream) {
59+
LaunchFill<Eigen::half>(out, value, n, stream);
60+
}
61+
62+
void LaunchMusaFill_bfloat16(Eigen::bfloat16* out, Eigen::bfloat16 value,
63+
int64_t n, musaStream_t stream) {
64+
LaunchFill<Eigen::bfloat16>(out, value, n, stream);
65+
}
66+
67+
void LaunchMusaFill_bool(bool* out, bool value, int64_t n, musaStream_t stream) {
68+
LaunchFill<bool>(out, value, n, stream);
69+
}
70+
71+
} // extern "C"
72+
73+
} // namespace musa
74+
} // namespace tensorflow

‎musa_ext/kernels/array/musa_fill_op.cc‎

Lines changed: 120 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
#include "../utils_op.h"
2-
#include "musa_fill_functor.h"
2+
#include "mu/device/musa_memcpy.h"
3+
4+
#include <algorithm>
5+
#include <musa_runtime.h>
6+
#include <vector>
7+
38
#include "tensorflow/core/framework/op_kernel.h"
49
#include "tensorflow/core/framework/register_types.h"
510
#include "tensorflow/core/framework/tensor.h"
@@ -12,6 +17,22 @@ namespace musa {
1217

1318
namespace {
1419

20+
extern "C" {
21+
void LaunchMusaFill_float(float* out, float value, int64_t n,
22+
musaStream_t stream);
23+
void LaunchMusaFill_double(double* out, double value, int64_t n,
24+
musaStream_t stream);
25+
void LaunchMusaFill_int32(int32* out, int32 value, int64_t n,
26+
musaStream_t stream);
27+
void LaunchMusaFill_int64(int64* out, int64 value, int64_t n,
28+
musaStream_t stream);
29+
void LaunchMusaFill_half(Eigen::half* out, Eigen::half value, int64_t n,
30+
musaStream_t stream);
31+
void LaunchMusaFill_bfloat16(Eigen::bfloat16* out, Eigen::bfloat16 value,
32+
int64_t n, musaStream_t stream);
33+
void LaunchMusaFill_bool(bool* out, bool value, int64_t n, musaStream_t stream);
34+
}
35+
1536
template <typename T, typename... Rest>
1637
struct is_any : std::false_type {};
1738

@@ -23,31 +44,88 @@ struct is_any<T, First, Rest...>
2344
: std::integral_constant<bool, std::is_same<T, First>::value ||
2445
is_any<T, Rest...>::value> {};
2546

26-
// template <typename T>
27-
// Status MusaFillCall(Tensor* out, T value, OpKernelContext* context) {
28-
// mFill op;
29-
// mHandle& h = GetHandleByCtx(context);
30-
// auto out_mt = CreateMTensor(*out);
31-
32-
// if (is_any<T, int8, int16, int, int64, uint8, uint16, uint32, uint64,
33-
// bool>::value) {
34-
// if (mStatus::SUCCESS != op.SetValue(static_cast<int64_t>(value))) {
35-
// return errors::Internal("mtdnn set value (int) error!");
36-
// }
37-
// } else if (is_any<T, float, double, Eigen::half, Eigen::bfloat16>::value) {
38-
// if (mStatus::SUCCESS != op.SetValue(static_cast<double>(value))) {
39-
// return errors::Internal("mtdnn set value (float) error!");
40-
// }
41-
// } else {
42-
// return errors::Unimplemented("Data type not supported in MTGPU Fill.");
43-
// }
44-
45-
// if (mStatus::SUCCESS != op.Run(h, out_mt)) {
46-
// return errors::Internal("mtdnn run op error!");
47-
// }
48-
49-
// return OkStatus();
50-
// }
47+
template <typename T>
48+
Status LaunchFillKernel(T* out, T value, int64_t n, musaStream_t stream);
49+
50+
#define DEFINE_FILL_LAUNCHER(T, suffix) \
51+
template <> \
52+
Status LaunchFillKernel<T>(T* out, T value, int64_t n, musaStream_t stream) { \
53+
LaunchMusaFill_##suffix(out, value, n, stream); \
54+
musaError_t err = musaGetLastError(); \
55+
if (err != musaSuccess) { \
56+
return errors::Internal("MUSA Fill kernel launch failed: ", \
57+
musaGetErrorString(err)); \
58+
} \
59+
return OkStatus(); \
60+
}
61+
62+
DEFINE_FILL_LAUNCHER(float, float)
63+
DEFINE_FILL_LAUNCHER(double, double)
64+
DEFINE_FILL_LAUNCHER(int32, int32)
65+
DEFINE_FILL_LAUNCHER(int64, int64)
66+
DEFINE_FILL_LAUNCHER(Eigen::half, half)
67+
DEFINE_FILL_LAUNCHER(Eigen::bfloat16, bfloat16)
68+
DEFINE_FILL_LAUNCHER(bool, bool)
69+
70+
#undef DEFINE_FILL_LAUNCHER
71+
72+
bool IsDeviceReadablePointer(const void* ptr) {
73+
musaPointerAttributes attributes;
74+
musaError_t attr_err = musaPointerGetAttributes(&attributes, ptr);
75+
if (attr_err == musaSuccess) {
76+
return attributes.type == musaMemoryTypeDevice ||
77+
attributes.type == musaMemoryTypeManaged;
78+
}
79+
80+
// Pageable host memory is reported as unregistered on some MUSA versions.
81+
musaGetLastError();
82+
return false;
83+
}
84+
85+
template <typename T>
86+
Status ReadVectorInputToHost(OpKernelContext* context, int input_index,
87+
std::vector<T>* values) {
88+
const Tensor& tensor = context->input(input_index);
89+
values->resize(tensor.NumElements());
90+
if (values->empty()) {
91+
return OkStatus();
92+
}
93+
94+
const T* input_data = tensor.flat<T>().data();
95+
if (IsDeviceReadablePointer(input_data)) {
96+
mStatus copy_status =
97+
MusaMemcpyD2H(values->data(), input_data, tensor.TotalBytes());
98+
if (copy_status != mStatus::SUCCESS) {
99+
return errors::Internal("MUSA Fill failed to copy input ", input_index,
100+
" to host.");
101+
}
102+
} else {
103+
std::copy(input_data, input_data + values->size(), values->data());
104+
}
105+
return OkStatus();
106+
}
107+
108+
template <typename T>
109+
Status ReadScalarInputToHost(OpKernelContext* context, int input_index,
110+
T* value) {
111+
const Tensor& tensor = context->input(input_index);
112+
if (tensor.NumElements() < 1) {
113+
return errors::InvalidArgument("MUSA Fill input ", input_index,
114+
" must contain a scalar value.");
115+
}
116+
117+
const T* input_data = tensor.flat<T>().data();
118+
if (IsDeviceReadablePointer(input_data)) {
119+
mStatus copy_status = MusaMemcpyD2H(value, input_data, sizeof(T));
120+
if (copy_status != mStatus::SUCCESS) {
121+
return errors::Internal("MUSA Fill failed to copy scalar input ",
122+
input_index, " to host.");
123+
}
124+
} else {
125+
*value = input_data[0];
126+
}
127+
return OkStatus();
128+
}
51129

52130
} // namespace
53131

@@ -78,21 +156,29 @@ class MusaFillOp : public MusaOpKernel {
78156
errors::InvalidArgument("value must represent a scalar, got shape ",
79157
Tvalue.shape().DebugString()));
80158

81-
auto dims_vec = Tdims.flat<Index>();
159+
std::vector<Index> host_dims;
160+
OP_REQUIRES_OK(context, ReadVectorInputToHost<Index>(context, 0, &host_dims));
161+
82162
TensorShape shape;
83-
OP_REQUIRES_OK(context, TensorShapeUtils::MakeShape(
84-
reinterpret_cast<const Index*>(dims_vec.data()),
85-
dims_vec.size(), &shape));
163+
OP_REQUIRES_OK(
164+
context,
165+
TensorShapeUtils::MakeShape(host_dims.data(), host_dims.size(), &shape));
86166

87167
Tensor* out = nullptr;
88168
OP_REQUIRES_OK(context, context->allocate_output(0, shape, &out));
89169

90170
if (shape.num_elements() == 0) return;
91171

92-
auto out_mt = CreateMTensor(*out);
93-
OP_REQUIRES_OK(
94-
context,
95-
MusaFillCall(&out_mt, static_cast<T*>(Tvalue.data())[0], context));
172+
musaStream_t stream = GetMusaStreamByCtx(context);
173+
OP_REQUIRES(context, stream != nullptr,
174+
errors::Internal("MUSA stream is unavailable for Fill"));
175+
176+
T host_value;
177+
OP_REQUIRES_OK(context, ReadScalarInputToHost<T>(context, 1, &host_value));
178+
179+
OP_REQUIRES_OK(context,
180+
LaunchFillKernel<T>(out->flat<T>().data(), host_value,
181+
shape.num_elements(), stream));
96182
}
97183
};
98184

0 commit comments

Comments
 (0)