Skip to content

Commit fd0cd6a

Browse files
naromero77amdpytorchmergebot
naromero77amd
authored and
pytorchmergebot
committed
[ROCm][TunableOp] Improve identification of fastest solution (pytorch#144942)
This PR addresses some stability issues with identifying the fastest solution on AMD GPUs, particularly the MI300. Changes include: - An improved timer, StreamTimerNoSync - More aggressive skipping of slow solutions - Additional statistics that can be used for diagnostics PYTORCH_TUNABLEOP_VERBOSE=3 Pull Request resolved: pytorch#144942 Approved by: https://github.com/jeffdaily
1 parent e20b0c8 commit fd0cd6a

File tree

6 files changed

+172
-26
lines changed

6 files changed

+172
-26
lines changed

‎aten/src/ATen/cuda/tunable/GemmHipblaslt.h

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -598,13 +598,6 @@ auto GetHipBlasLtTypeStringAndOps() {
598598
heuristic_result));
599599
TORCH_HIPBLASLT_CHECK(hipblasLtDestroy(handle));
600600

601-
// Sort heuristic_result by algo index to make sure the order of returned algos is deterministic.
602-
std::sort(heuristic_result.begin(),
603-
heuristic_result.end(),
604-
[](hipblasLtMatmulHeuristicResult_t& a, hipblasLtMatmulHeuristicResult_t& b) {
605-
return hipblaslt_ext::getIndexFromAlgo(a.algo) < hipblaslt_ext::getIndexFromAlgo(b.algo);
606-
});
607-
608601
int returned_algo_count = heuristic_result.size();
609602
std::vector<std::pair<std::string, std::unique_ptr<Callable<ParamsT>>>> ret;
610603
for (int i = 0; i < returned_algo_count; i++) {

‎aten/src/ATen/cuda/tunable/GemmRocblas.h

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -192,9 +192,6 @@ auto GetRocBlasGemmTypeStringAndOps() {
192192
rocblas_gemm_flags_none,
193193
solutions.data(),
194194
&solution_size));
195-
// Sort the solutions in ascending order to make the solution vector deterministic across runs
196-
std::sort(solutions.begin(), solutions.end());
197-
198195
std::vector<std::pair<std::string, std::unique_ptr<Callable<GemmParams<T>>>>> ret;
199196
for (size_t i = 0; i < solutions.size(); ++i) {
200197
auto callable = std::make_unique<RocblasGemmOp<T>>(solutions[i]);

‎aten/src/ATen/cuda/tunable/StreamTimer.cpp

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ StreamTimer::StreamTimer() {
2424
StreamTimer::~StreamTimer() = default;
2525

2626
void StreamTimer::Start() {
27-
AT_CUDA_CHECK(cudaDeviceSynchronize());
27+
AT_CUDA_CHECK(cudaEventSynchronize(start_));
2828
AT_CUDA_CHECK(cudaEventRecord(start_, at::cuda::getCurrentCUDAStream()));
2929
}
3030

@@ -40,4 +40,27 @@ float StreamTimer::Duration() {
4040
return time;
4141
}
4242

43+
StreamTimerNoSync::StreamTimerNoSync() {
44+
AT_CUDA_CHECK(cudaEventCreate(&start_));
45+
AT_CUDA_CHECK(cudaEventCreate(&end_));
46+
}
47+
48+
StreamTimerNoSync::~StreamTimerNoSync() = default;
49+
50+
void StreamTimerNoSync::Start() {
51+
AT_CUDA_CHECK(cudaEventRecord(start_, at::cuda::getCurrentCUDAStream()));
52+
}
53+
54+
void StreamTimerNoSync::End() {
55+
AT_CUDA_CHECK(cudaEventRecord(end_, at::cuda::getCurrentCUDAStream()));
56+
}
57+
58+
float StreamTimerNoSync::Duration() {
59+
auto time = std::numeric_limits<float>::quiet_NaN();
60+
AT_CUDA_CHECK(cudaEventSynchronize(end_));
61+
// time is in ms with a resolution of 1 us
62+
AT_CUDA_CHECK(cudaEventElapsedTime(&time, start_, end_));
63+
return time;
64+
}
65+
4366
} // namespace at::cuda::tunable

‎aten/src/ATen/cuda/tunable/StreamTimer.h

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,20 @@ class StreamTimer : public ITimer {
3131
cudaEvent_t end_{};
3232
};
3333

34+
class StreamTimerNoSync : public ITimer {
35+
public:
36+
StreamTimerNoSync();
37+
~StreamTimerNoSync() override;
38+
39+
void Start() override;
40+
41+
void End() override;
42+
43+
float Duration() override;
44+
45+
private:
46+
cudaEvent_t start_{};
47+
cudaEvent_t end_{};
48+
};
49+
3450
} // namespace at::cuda::tunable

‎aten/src/ATen/cuda/tunable/TunableGemm.h

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
#include <ATen/cuda/tunable/GemmHipblaslt.h>
1515
#include <ATen/cuda/tunable/GemmRocblas.h>
1616
#endif
17-
#include <ATen/cuda/tunable/StreamTimer.h>
1817
#include <ATen/cuda/tunable/TunableOp.h>
1918
#include <c10/cuda/CUDACachingAllocator.h>
2019
#include <c10/util/Float8_e4m3fn.h>
@@ -192,7 +191,7 @@ inline const char* TypeName(c10::complex<float> v) {
192191
}
193192

194193
template <typename T, BlasOp ALayout, BlasOp BLayout>
195-
class GemmTunableOp : public TunableOp<GemmParams<T>, StreamTimer> {
194+
class GemmTunableOp : public TunableOp<GemmParams<T>> {
196195
public:
197196
GemmTunableOp() {
198197
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmOp<T>>());
@@ -217,6 +216,8 @@ class GemmTunableOp : public TunableOp<GemmParams<T>, StreamTimer> {
217216
}
218217
}
219218
#endif
219+
220+
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmOp<T>>());
220221
}
221222

222223
std::string Signature() override {
@@ -225,7 +226,7 @@ class GemmTunableOp : public TunableOp<GemmParams<T>, StreamTimer> {
225226
};
226227

227228
template <typename T, BlasOp ALayout, BlasOp BLayout>
228-
class GemmAndBiasTunableOp : public TunableOp<GemmAndBiasParams<T>, StreamTimer> {
229+
class GemmAndBiasTunableOp : public TunableOp<GemmAndBiasParams<T>> {
229230
public:
230231
GemmAndBiasTunableOp() {
231232
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmAndBiasOp<T>>());
@@ -243,6 +244,8 @@ class GemmAndBiasTunableOp : public TunableOp<GemmAndBiasParams<T>, StreamTimer>
243244
}
244245
}
245246
#endif
247+
248+
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmAndBiasOp<T>>());
246249
}
247250

248251
std::string Signature() override {
@@ -251,7 +254,7 @@ class GemmAndBiasTunableOp : public TunableOp<GemmAndBiasParams<T>, StreamTimer>
251254
};
252255

253256
template <typename T, BlasOp ALayout, BlasOp BLayout>
254-
class GemmStridedBatchedTunableOp : public TunableOp<GemmStridedBatchedParams<T>, StreamTimer> {
257+
class GemmStridedBatchedTunableOp : public TunableOp<GemmStridedBatchedParams<T>> {
255258
public:
256259
GemmStridedBatchedTunableOp() {
257260
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmStridedBatchedOp<T>>());
@@ -276,6 +279,8 @@ class GemmStridedBatchedTunableOp : public TunableOp<GemmStridedBatchedParams<T>
276279
}
277280
}
278281
#endif
282+
283+
this->RegisterOp(std::string("Default"), std::make_unique<DefaultGemmStridedBatchedOp<T>>());
279284
}
280285

281286
std::string Signature() override {
@@ -284,7 +289,7 @@ class GemmStridedBatchedTunableOp : public TunableOp<GemmStridedBatchedParams<T>
284289
};
285290

286291
template <typename AT, typename BT, typename CT, BlasOp ALayout, BlasOp BLayout>
287-
class ScaledGemmTunableOp : public TunableOp<ScaledGemmParams<CT>, StreamTimer> {
292+
class ScaledGemmTunableOp : public TunableOp<ScaledGemmParams<CT>> {
288293
public:
289294
ScaledGemmTunableOp() {
290295
this->RegisterOp(std::string("Default"), std::make_unique<DefaultScaledGemmOp<CT>>());
@@ -294,6 +299,8 @@ class ScaledGemmTunableOp : public TunableOp<ScaledGemmParams<CT>, StreamTimer>
294299
this->RegisterOp(std::move(name), std::move(op));
295300
}
296301
#endif
302+
303+
this->RegisterOp(std::string("Default"), std::make_unique<DefaultScaledGemmOp<CT>>());
297304
}
298305

299306
std::string Signature() override {

‎aten/src/ATen/cuda/tunable/TunableOp.h

Lines changed: 120 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#pragma once
1111

1212
#include <ATen/cuda/tunable/Tunable.h>
13+
#include <ATen/cuda/tunable/StreamTimer.h>
1314
#include <ATen/cuda/Sleep.h>
1415
#include <c10/cuda/CUDACachingAllocator.h>
1516

@@ -35,7 +36,57 @@ class Callable {
3536
}
3637
};
3738

38-
template <typename ParamsT, typename TimerT>
39+
namespace {
40+
41+
/** http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance */
42+
43+
class Stats {
44+
public:
45+
Stats() {
46+
_n = 0UL;
47+
_mean = 0.0;
48+
_M2 = 0.0;
49+
_sum = 0.0;
50+
_min = 0.0;
51+
_max = 0.0;
52+
}
53+
54+
void sample_value(const double x) {
55+
double delta = 0;
56+
_sum = _sum + x;
57+
if (0UL == _n) {
58+
_min = x;
59+
_max = x;
60+
}
61+
else {
62+
_min = _min < x ? _min : x;
63+
_max = _max > x ? _max : x;
64+
}
65+
_n = _n + 1UL;
66+
delta = x - _mean;
67+
_mean = _mean + delta/_n;
68+
_M2 = _M2 + delta * (x - _mean);
69+
}
70+
71+
double variance() const {
72+
return _M2/(_n-1);
73+
}
74+
75+
double stddev() const {
76+
return std::sqrt(variance());
77+
}
78+
79+
unsigned long _n;
80+
double _mean;
81+
double _M2;
82+
double _sum;
83+
double _min;
84+
double _max;
85+
};
86+
87+
} // anonymous namespace
88+
89+
template <typename ParamsT>
3990
class TunableOp {
4091
public:
4192
virtual ~TunableOp() = default;
@@ -100,10 +151,17 @@ class TunableOp {
100151
}
101152
}
102153

103-
static double Profile(Callable<ParamsT> *op, const std::vector<ParamsT*> &param, size_t num_iter, size_t &offset) {
154+
static double ProfileSimple(Callable<ParamsT> *op, const std::vector<ParamsT*> &param, size_t num_iter, size_t &offset) {
104155
TuningContext* ctx = getTuningContext();
105156
bool do_flush = ctx->IsICacheFlushEnabled();
106-
TimerT timer{};
157+
StreamTimerNoSync timer{};
158+
159+
// Small Mandatory Warmup
160+
// Reduces outliers
161+
for (size_t i = 0; i < 2; i++) {
162+
TORCH_CHECK(op->Call(param[(i+offset++)%param.size()]) == OK);
163+
}
164+
107165
timer.Start();
108166
for (size_t i = 0; i < num_iter; i++) {
109167
if (do_flush) {
@@ -115,6 +173,32 @@ class TunableOp {
115173
return timer.Duration() / num_iter;
116174
}
117175

176+
static Stats ProfileStats(Callable<ParamsT> *op, const std::vector<ParamsT*> &param, size_t num_iter, size_t &offset) {
177+
TuningContext* ctx = getTuningContext();
178+
bool do_flush = ctx->IsICacheFlushEnabled();
179+
std::vector<StreamTimerNoSync> timer(num_iter);
180+
181+
// Small Mandatory Warmup
182+
// Reduces outliers
183+
for (size_t i = 0; i < 2; i++) {
184+
TORCH_CHECK(op->Call(param[(i+offset++)%param.size()]) == OK);
185+
}
186+
187+
for (size_t i = 0; i < num_iter; i++) {
188+
timer[i].Start();
189+
TORCH_CHECK(op->Call(param[(i+offset++)%param.size()]) == OK);
190+
timer[i].End();
191+
if (do_flush) {
192+
at::cuda::flush_icache();
193+
}
194+
}
195+
Stats s;
196+
for (size_t i = 0; i < num_iter; i++) {
197+
s.sample_value(timer[i].Duration());
198+
}
199+
return s;
200+
}
201+
118202
protected:
119203
virtual ResultEntry FindFastest(const ParamsT* params) {
120204
TuningContext* ctx = getTuningContext();
@@ -184,14 +268,25 @@ class TunableOp {
184268
}
185269

186270
// collect a small profile
187-
constexpr const int approx_num_iter = 3;
188-
auto approx_duration = Profile(candidate, reusable_params, approx_num_iter, offset);
271+
int approx_num_iter = 3;
272+
auto s = ProfileStats(candidate, reusable_params, approx_num_iter, offset);
273+
double approx_duration = s._mean;
189274
// bail if too slow
190-
if (approx_duration > 2 * min_duration_ms) {
275+
if (approx_duration > 1.5 * min_duration_ms) {
191276
TUNABLE_LOG3("├──skip slow instance id=", i, ", ", op_sig, '(', params_sig, ") ", op_names_[i]);
192277
continue;
193278
}
194279

280+
// 2nd phase skip, more aggressive
281+
approx_num_iter = 10;
282+
s = ProfileStats(candidate, reusable_params, approx_num_iter, offset);
283+
approx_duration = s._mean;
284+
// bail if too slow
285+
if (approx_duration > 1.15 * min_duration_ms) {
286+
TUNABLE_LOG3("├──2nd skip slow instance id=", i, ", ", op_sig, '(', params_sig, ") ", op_names_[i]);
287+
continue;
288+
}
289+
195290
// for warmup does user set max duration, max iters, or both?
196291
// warmup is allowed to be skipped by setting either iterations or duration to 0
197292
double max_warmup_duration = ctx->GetMaxWarmupDurationMs();
@@ -238,12 +333,27 @@ class TunableOp {
238333
"instance id=", i, ", ", op_sig, "(", params_sig, ") ", op_names_[i]);
239334
TUNABLE_LOG3("├──offset at ", offset);
240335
WarmUp(candidate, reusable_params, warmup_iter, offset);
241-
auto duration_ms = Profile(candidate, reusable_params, tuning_iter, offset);
242-
if (duration_ms < min_duration_ms) {
243-
TUNABLE_LOG3("├──found better instance id=", i, ". " , duration_ms, "ms. ", op_names_[i]);
244-
min_duration_ms = duration_ms;
336+
s = ProfileStats(candidate, reusable_params, tuning_iter, offset);
337+
auto s_stddev = s.stddev();
338+
// Assume normal distribution.
339+
// Solution with smallest mean + 2*sigma will be a better solution?
340+
// if ((s._mean + 2*s_stddev) < (min_duration_ms + 2*min_stddev_ms)) {
341+
if (s._mean < min_duration_ms) {
342+
TUNABLE_LOG3("├──found better instance id=", i, ". " , s._mean, "ms. ", op_names_[i],
343+
" min ", s._min,
344+
" max ", s._max,
345+
" mean ", s._mean,
346+
" std ", s_stddev);
347+
min_duration_ms = s._mean;
245348
id_name = op_names_[i];
246349
}
350+
else {
351+
TUNABLE_LOG3("├──found slower instance id=", i, ". " , s._mean, "ms. ", op_names_[i],
352+
" min ", s._min,
353+
" max ", s._max,
354+
" mean ", s._mean,
355+
" std ", s_stddev);
356+
}
247357
}
248358

249359
for (size_t i = 0; i < reusable_params.size(); i++) {

0 commit comments

Comments
 (0)