Add bfloat16 support to OCL bn kernels #3351

BrianHarrisonAMD · 2024-11-01T01:39:05Z

Draft of updating OCL BN to support bf16

src/kernels/MIOpenBatchNormFwdTrainSpatial.cl

bpepers-me · 2024-11-01T03:29:30Z

src/kernels/MIOpenBatchNormBwdSpatial.cl

-        buff[gammaindex]        = (_FLOAT_PREC)dscale;
-        buff[betaindex]         = (_FLOAT_PREC)dbias;
+        buff[gammaindex]        = FLOAT2FLOATPREC(dscale);
+        buff[betaindex]         = FLOAT2FLOATPREC(dbias);


this is odd since buff is _FLOAT*. but then a bunch of the functions around here don't seem right like the ones that only use _FLOAT so they better never be used for bf16.

Yea these look incorrect for BF16, and just generally, but I think these don't get called for mixed.
Only issue is it might have compilation errors that were okay before let me check how it behaved for FP16.

bpepers-me · 2024-11-01T03:30:04Z

src/kernels/MIOpenBatchNormBwdSpatial.cl

-        delta_scale[xgid] = (_FLOAT_PREC)ds;
-        delta_bias[xgid]  = (_FLOAT_PREC)db;
+        delta_scale[xgid] = FLOAT2FLOATPREC(ds);
+        delta_bias[xgid]  = FLOAT2FLOATPREC(db);


src/kernels/MIOpenBatchNormActivFwdTrainSpatial.cl

src/kernels/MIOpenBatchNormBwdPerAct.cl

src/kernels/MIOpenBatchNormActivFwdTrainSpatial.cl

src/kernels/batchnorm_functions.h

src/kernels/MIOpenBatchNormBwdPerAct.cl

…arriso/bn-ocl-bf16

* fix driver for old ocl kernel * fix template name * fix review comments * fix isApplicable for CK and OCL bn spatial * fix gtest issue in bn for ck and ocl * fix minor issue * fix clang format * hip tidy * fix gtest sample and fix hiptidy * * move type check logic to problem description of batch norm * add type checks to other ocl solvers * fix other minor issues * fix network cache * fix minor testing issue * disable ck bn for now --------- Co-authored-by: BrianHarrisonAMD <[email protected]>

averinevg

Post-merge review

@BrianHarrisonAMD Please create post-merge fix.

CC @junliume @bpepers-me

averinevg · 2024-11-04T21:10:18Z

src/batchnorm/problem_description.cpp

+bool is_fp16_or_bfp16(miopenDataType_t type)
+{
+    return ((type == miopenHalf) || (type == miopenBFloat16));
+}
+
+bool is_fp32_or_fp64(miopenDataType_t type)
+{
+    return ((type == miopenFloat) || (type == miopenDouble));
+}


Please don't use snake case for function names

averinevg · 2024-11-04T21:17:25Z

src/solver/batchnorm/forward_spatial_single.cpp

+    const ExecutionContext&, const miopen::batchnorm::ProblemDescription& bn_problem) const
 {
-    if(problem.GetDirection() != miopen::batchnorm::Direction::ForwardTraining ||
-       problem.GetMode() != miopenBNSpatial)
+    if(bn_problem.GetDirection() != miopen::batchnorm::Direction::ForwardTraining ||
+       bn_problem.GetMode() != miopenBNSpatial)


This renaming was not superfluous. I believe that the name problem in this context is sufficient. Inside all solvers this structure is called exactly this way, and renaming it will only bring chaos.

averinevg · 2024-11-04T21:18:12Z

src/batchnorm/problem_description.cpp

+bool IsOCLInferTypeValid(const ProblemDescription& bn_problem)
+{
+    // case 1 : mix type
+    return (
+        (is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+         is_fp16_or_bfp16(bn_problem.GetYDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnBias().GetType())) ||
+        // case 2 : float type
+        (is_fp32(bn_problem.GetXDesc().GetType()) && is_fp32(bn_problem.GetYDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnBias().GetType())));
+}
+
+bool IsCKInferTypeValid(const ProblemDescription& bn_problem)
+{
+    // case 1 : mix type
+    return ((is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetYDesc().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetBnScale().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetBnBias().GetType()) &&
+             is_fp32(bn_problem.GetBnSMean().GetType()) &&
+             is_fp32(bn_problem.GetBnSVar().GetType())) ||
+            // case 2 : fp32 or fp64
+            (is_fp32_or_fp64(bn_problem.GetXDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetYDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnScale().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnBias().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSMean().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSVar().GetType())));
+}
+
+bool IsOCLFwdTrainTypeValid(const ProblemDescription& bn_problem)
+{
+    // case 1 : mix type
+    return (
+        (is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+         is_fp16_or_bfp16(bn_problem.GetYDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnBias().GetType())) ||
+        // case 2 : float type
+        (is_fp32(bn_problem.GetXDesc().GetType()) && is_fp32(bn_problem.GetYDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnBias().GetType())));
+}
+
+bool IsCKFwdTrainTypeValid(const ProblemDescription& bn_problem)
+{
+    // case 1 : mix type
+    return ((is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetYDesc().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetBnScale().GetType()) &&
+             is_fp16_or_bfp16(bn_problem.GetBnBias().GetType()) &&
+             is_fp32(bn_problem.GetBnSMean().GetType()) &&
+             is_fp32(bn_problem.GetBnSVar().GetType())) ||
+            // case 2 : fp32 or fp64
+            (is_fp32_or_fp64(bn_problem.GetXDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetYDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnScale().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnBias().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSMean().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSVar().GetType())));
+}
+
+bool IsOCLBwdTypeValid(const ProblemDescription& bn_problem)
+{
+    return (
+        (is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+         is_fp16_or_bfp16(bn_problem.GetDXDesc().GetType()) &&
+         is_fp16_or_bfp16(bn_problem.GetDYDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnSMean().GetType()) &&
+         is_fp32(bn_problem.GetBnSVar().GetType())) ||
+        // case 1 : fp32
+        (is_fp32(bn_problem.GetXDesc().GetType()) && is_fp32(bn_problem.GetDXDesc().GetType()) &&
+         is_fp32(bn_problem.GetBnScale().GetType()) && is_fp32(bn_problem.GetBnBias().GetType()) &&
+         is_fp32(bn_problem.GetBnSMean().GetType()) && is_fp32(bn_problem.GetBnSVar().GetType())));
+}
+
+bool IsCKBwdTypeValid(const ProblemDescription& bn_problem)
+{
+    return ((is_fp16_or_bfp16(bn_problem.GetXDesc().GetType()) &&
+             bn_problem.GetDXDesc().GetType() == miopenFloat &&
+             is_fp16_or_bfp16(bn_problem.GetBnScale().GetType()) &&
+             bn_problem.GetDYDesc().GetType() == miopenFloat &&
+             bn_problem.GetBnSMean().GetType() == miopenFloat &&
+             bn_problem.GetBnSVar().GetType() == miopenFloat) ||
+            // case 1 : fp32 or fp64
+            (is_fp32_or_fp64(bn_problem.GetXDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetDXDesc().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnScale().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnBias().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSMean().GetType()) &&
+             is_fp32_or_fp64(bn_problem.GetBnSVar().GetType())));
+}


These functions should not be here. ProblemDescription is a universal structure, not tied to a backend, third-party library, etc. Move them to the appropriate files.

Add bfloat16 support to OCL bn kernels

e44dd04

bpepers-me reviewed Nov 1, 2024

View reviewed changes

src/kernels/MIOpenBatchNormFwdTrainSpatial.cl Outdated Show resolved Hide resolved

BrianHarrisonAMD added 3 commits November 1, 2024 02:04

Fix issue

c016604

Fix spacing

a9e0ecb

Add bf16 support for bwd OCL spatial

43aa479

bpepers-me reviewed Nov 1, 2024

View reviewed changes

Further changes of OCL kernels for bfp16

ce33a3e

BrianHarrisonAMD commented Nov 1, 2024

View reviewed changes

src/kernels/MIOpenBatchNormActivFwdTrainSpatial.cl Outdated Show resolved Hide resolved

BrianHarrisonAMD commented Nov 1, 2024

View reviewed changes

src/kernels/MIOpenBatchNormBwdPerAct.cl Show resolved Hide resolved

BrianHarrisonAMD added 4 commits November 1, 2024 14:16

Add bf16 support for remaining OCL solvers

09de609

Fix compilation error

f0a080c

Fix compilation errors

8acd52d

Fix formatting

24b0a11

CAHEK7 suggested changes Nov 1, 2024

View reviewed changes

src/kernels/MIOpenBatchNormActivFwdTrainSpatial.cl Outdated Show resolved Hide resolved

src/kernels/batchnorm_functions.h Outdated Show resolved Hide resolved

src/kernels/batchnorm_functions.h Show resolved Hide resolved

BrianHarrisonAMD added 4 commits November 1, 2024 15:48

Fix extra bracket being added

c7bc908

Fix kernel compilation and issue with variant selection for BF16 mixed

8ff2396

Swap solver order to prefer OCL over CK by default

99326bf

Update batchnorm bfp16mix to match fp16mix declaration

58bd202

BrianHarrisonAMD commented Nov 1, 2024

View reviewed changes

src/kernels/MIOpenBatchNormBwdPerAct.cl Outdated Show resolved Hide resolved

BrianHarrisonAMD added 3 commits November 1, 2024 18:18

Fix improper cast

90b0695

Fix formatting

59ad519

Merge remote-tracking branch 'origin/develop' into bharriso/bn-ocl-bf16

90237e0

BrianHarrisonAMD marked this pull request as ready for review November 1, 2024 19:33

BrianHarrisonAMD requested review from JehandadKhan and junliume as code owners November 1, 2024 19:33

BrianHarrisonAMD and others added 4 commits November 2, 2024 02:28

Merge develop

fffc1d1

Merge remote-tracking branch 'origin/develop' into bharriso/bn-ocl-bf16

d811667

Merge branch 'develop' into bharriso/bn-ocl-bf16

6e85cb9

Fix kernel typo for activ

a95531e

BrianHarrisonAMD and others added 3 commits November 2, 2024 19:04

Merge branch 'bharriso/bn-ocl-bf16' of github.com:ROCm/MIOpen into bh…

a792d50

…arriso/bn-ocl-bf16

Fix formatting

3d5cc27

junliume approved these changes Nov 2, 2024

View reviewed changes

bpepers-me approved these changes Nov 2, 2024

View reviewed changes

BrianHarrisonAMD merged commit 1922f62 into develop Nov 2, 2024
13 of 128 checks passed

BrianHarrisonAMD deleted the bharriso/bn-ocl-bf16 branch November 2, 2024 20:27

averinevg reviewed Nov 4, 2024

View reviewed changes

averinevg mentioned this pull request Nov 4, 2024

Fix post-merge comments in #3351 #3366

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add bfloat16 support to OCL bn kernels #3351

Add bfloat16 support to OCL bn kernels #3351

BrianHarrisonAMD commented Nov 1, 2024

bpepers-me Nov 1, 2024

BrianHarrisonAMD Nov 1, 2024 •

edited

Loading

bpepers-me Nov 1, 2024

averinevg left a comment

averinevg Nov 4, 2024

averinevg Nov 4, 2024

averinevg Nov 4, 2024

Add bfloat16 support to OCL bn kernels #3351

Add bfloat16 support to OCL bn kernels #3351

Conversation

BrianHarrisonAMD commented Nov 1, 2024

bpepers-me Nov 1, 2024

Choose a reason for hiding this comment

BrianHarrisonAMD Nov 1, 2024 • edited Loading

Choose a reason for hiding this comment

bpepers-me Nov 1, 2024

Choose a reason for hiding this comment

averinevg left a comment

Choose a reason for hiding this comment

averinevg Nov 4, 2024

Choose a reason for hiding this comment

averinevg Nov 4, 2024

Choose a reason for hiding this comment

averinevg Nov 4, 2024

Choose a reason for hiding this comment

BrianHarrisonAMD Nov 1, 2024 •

edited

Loading