Skip to content

[ExecuTorch] support BF16 in op_copy #4979

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
merged 17 commits into from
Sep 6, 2024
Merged

[ExecuTorch] support BF16 in op_copy #4979

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
4fb8d62
[ExecuTorch] Implement BFloat16 and hook it up to scalar_type_util
swolchok Aug 29, 2024
4e8b86b
[ExecuTorch] support BF16 in op_to_copy
swolchok Aug 29, 2024
ff88974
[ExecuTorch] support BF16 in op_mul
swolchok Aug 29, 2024
46579e4
[ExecuTorch] support BF16 in op_mm
swolchok Aug 29, 2024
8388e56
[ExecuTorch] support BF16 in op_copy
swolchok Aug 29, 2024
c2baa70
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Aug 30, 2024
1b9450a
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Aug 30, 2024
f54b32a
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Aug 30, 2024
e46249e
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Aug 30, 2024
7a3174d
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 3, 2024
a32cd02
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 3, 2024
90f5bd7
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 3, 2024
31f8690
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 3, 2024
51a8abb
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 4, 2024
44c4a2c
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 4, 2024
01741c9
Update base for Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 6, 2024
5a83963
Update on "[ExecuTorch] support BF16 in op_copy"
swolchok Sep 6, 2024
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
8 changes: 4 additions & 4 deletions kernels/portable/cpu/op_copy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ Tensor& copy_out(
ScalarType in_type = in.scalar_type();
ScalarType src_type = src.scalar_type();

ET_SWITCH_REALHB_TYPES(in_type, ctx, "copy.out", CTYPE, [&]() {
ET_SWITCH_REALHB_TYPES(src_type, ctx, "copy.out", CTYPE_SRC, [&]() {
ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, "copy.out", CTYPE, [&]() {
ET_SWITCH_REALHBBF16_TYPES(src_type, ctx, "copy.out", CTYPE_SRC, [&]() {
apply_binary_elementwise_fn<CTYPE, CTYPE_SRC, CTYPE>(
[](const CTYPE val_in, const CTYPE_SRC val_src) {
return convert<CTYPE, CTYPE_SRC>(val_src);
Expand Down Expand Up @@ -75,8 +75,8 @@ copy_(RuntimeContext& ctx, Tensor& in, const Tensor& src, bool non_blocking) {
ScalarType in_type = in.scalar_type();
ScalarType src_type = src.scalar_type();

ET_SWITCH_REALHB_TYPES(in_type, ctx, "copy_", CTYPE, [&]() {
ET_SWITCH_REALHB_TYPES(src_type, ctx, "copy_", CTYPE_SRC, [&]() {
ET_SWITCH_REALHBBF16_TYPES(in_type, ctx, "copy_", CTYPE, [&]() {
ET_SWITCH_REALHBBF16_TYPES(src_type, ctx, "copy_", CTYPE_SRC, [&]() {
apply_binary_elementwise_fn<CTYPE, CTYPE_SRC, CTYPE>(
[](const CTYPE val_in, const CTYPE_SRC val_src) {
return convert<CTYPE, CTYPE_SRC>(val_src);
Expand Down
27 changes: 14 additions & 13 deletions kernels/portable/cpu/op_mm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,19 +34,20 @@ mm_out(RuntimeContext& ctx, const Tensor& in, const Tensor& mat2, Tensor& out) {

ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);

ET_SWITCH_REAL_TYPES_AND(Half, in.scalar_type(), ctx, "mm.out", CTYPE, [&]() {
size_t m = in.size(0);
size_t n = in.size(1);
size_t p = mat2.size(1);

vec_matmul<CTYPE>(
out.mutable_data_ptr<CTYPE>(),
in.const_data_ptr<CTYPE>(),
mat2.const_data_ptr<CTYPE>(),
m,
n,
p);
});
ET_SWITCH_REAL_TYPES_AND2(
Half, BFloat16, in.scalar_type(), ctx, "mm.out", CTYPE, [&]() {
size_t m = in.size(0);
size_t n = in.size(1);
size_t p = mat2.size(1);

vec_matmul<CTYPE>(
out.mutable_data_ptr<CTYPE>(),
in.const_data_ptr<CTYPE>(),
mat2.const_data_ptr<CTYPE>(),
m,
n,
p);
});

return out;
}
Expand Down
4 changes: 2 additions & 2 deletions kernels/test/op_copy_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,13 +125,13 @@ class OpCopyInplaceTest : public OperatorTest {
// regular test for copy.out
TEST_F(OpCopyTest, AllRealDtypesSupported) {
#define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>();
ET_FORALL_REAL_TYPES(TEST_ENTRY);
ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
#undef TEST_ENTRY
}

TEST_F(OpCopyTest, EmptyInputSupported) {
#define TEST_ENTRY(ctype, dtype) test_empty_input<ctype, ScalarType::dtype>();
ET_FORALL_REAL_TYPES_AND(Bool, TEST_ENTRY);
ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
#undef TEST_ENTRY
}

Expand Down
2 changes: 1 addition & 1 deletion kernels/test/op_mm_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ TEST_F(OpMmOutTest, OutputDim) {
/// zeros().
TEST_F(OpMmOutTest, AllDtypesSupported) {
#define TEST_ENTRY(ctype, dtype) test_dtype<ctype, ScalarType::dtype>();
ET_FORALL_REAL_TYPES_AND(Half, TEST_ENTRY);
ET_FORALL_REALHBF16_TYPES(TEST_ENTRY);
#undef TEST_ENTRY
// TODO: Also add tests for half, complex, quantized, and other types. Easiest
// way to do that would be to make TensorFactory support zeros() and ones()
Expand Down
Loading