Skip to content

Commit

Permalink
Add pattern template for binary bitwise ops (#3486)
Browse files Browse the repository at this point in the history 
Summary:
Pull Request resolved: #3486

Similar to D56744651.

Differential Revision: D56852163
  • Loading branch information
@swolchok @facebook-github-bot
swolchok authored and facebook-github-bot committed May 3, 2024
1 parent e0dadf0 commit 1693060
Show file tree
Hide file tree
Showing 6 changed files with 95 additions and 212 deletions.
77 changes: 6 additions & 71 deletions kernels/portable/cpu/op_bitwise_and.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@
* LICENSE file in the root directory of this source tree.
*/

#include <cmath>
// patternlint-disable-next-line executorch-cpp-nostdinc
#include <functional>

#include <executorch/kernels/portable/cpu/pattern/bitwise_op.h>
#include <executorch/kernels/portable/cpu/scalar_utils.h>
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
#include <executorch/kernels/portable/cpu/util/functional_util.h>
Expand All @@ -17,81 +19,14 @@ namespace torch {
namespace executor {
namespace native {

namespace {

template <typename CTYPE>
CTYPE bitwise_and(CTYPE a, CTYPE b) {
return a & b;
}

template <>
bool bitwise_and<bool>(bool a, bool b) {
return a && b;
}

} // namespace

using Tensor = exec_aten::Tensor;

Tensor& bitwise_and_Tensor_out(
RuntimeContext& ctx,
const Tensor& a,
const Tensor& b,
Tensor& out) {
ET_KERNEL_CHECK(
ctx,
resize_to_broadcast_target_size(a, b, out) == Error::Ok,
InvalidArgument,
out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = b.scalar_type();
ScalarType common_type = promoteTypes(a_type, b_type);
ScalarType out_type = out.scalar_type();

ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);

ET_SWITCH_INT_TYPES_AND(
Bool, a_type, ctx, "bitwise_and.Tensor_out", CTYPE_A, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool, b_type, ctx, "bitwise_and.Tensor_out", CTYPE_B, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool,
common_type,
ctx,
"bitwise_and.Tensor_out",
CTYPE_IN,
[&]() {
ET_SWITCH_REAL_TYPES_AND(
Bool,
out_type,
ctx,
"bitwise_and.Tensor_out",
CTYPE_OUT,
[&]() {
apply_binary_elementwise_fn<
CTYPE_A,
CTYPE_B,
CTYPE_OUT>(
[](const CTYPE_A val_a, const CTYPE_B val_b) {
CTYPE_IN a_casted =
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value =
bitwise_and(a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
a,
b,
out);
});
});
});
});

return out;
return bitwise_op_out<std::bit_and>(ctx, a, b, out, "bitwise_and.Tensor_out");
}

Tensor& bitwise_and_Scalar_out(
Expand Down Expand Up @@ -142,8 +77,8 @@ Tensor& bitwise_and_Scalar_out(
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value =
bitwise_and(a_casted, b_casted);
CTYPE_IN value = std::bit_and<CTYPE_IN>()(
a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
Expand Down
75 changes: 6 additions & 69 deletions kernels/portable/cpu/op_bitwise_or.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@
* LICENSE file in the root directory of this source tree.
*/

#include <cmath>
// patternlint-disable-next-line executorch-cpp-nostdinc
#include <functional>

#include <executorch/kernels/portable/cpu/pattern/bitwise_op.h>
#include <executorch/kernels/portable/cpu/scalar_utils.h>
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
#include <executorch/kernels/portable/cpu/util/functional_util.h>
Expand All @@ -17,80 +19,14 @@ namespace torch {
namespace executor {
namespace native {

namespace {

template <typename CTYPE>
CTYPE bitwise_or(CTYPE a, CTYPE b) {
return a | b;
}

template <>
bool bitwise_or<bool>(bool a, bool b) {
return a || b;
}

} // namespace

using Tensor = exec_aten::Tensor;

Tensor& bitwise_or_Tensor_out(
RuntimeContext& ctx,
const Tensor& a,
const Tensor& b,
Tensor& out) {
ET_KERNEL_CHECK(
ctx,
resize_to_broadcast_target_size(a, b, out) == Error::Ok,
InvalidArgument,
out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = b.scalar_type();
ScalarType common_type = promoteTypes(a_type, b_type);
ScalarType out_type = out.scalar_type();

ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);

ET_SWITCH_INT_TYPES_AND(
Bool, a_type, ctx, "bitwise_or.Tensor_out", CTYPE_A, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool, b_type, ctx, "bitwise_or.Tensor_out", CTYPE_B, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool,
common_type,
ctx,
"bitwise_or.Tensor_out",
CTYPE_IN,
[&]() {
ET_SWITCH_REAL_TYPES_AND(
Bool,
out_type,
ctx,
"bitwise_or.Tensor_out",
CTYPE_OUT,
[&]() {
apply_binary_elementwise_fn<
CTYPE_A,
CTYPE_B,
CTYPE_OUT>(
[](const CTYPE_A val_a, const CTYPE_B val_b) {
CTYPE_IN a_casted =
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value = bitwise_or(a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
a,
b,
out);
});
});
});
});

return out;
return bitwise_op_out<std::bit_or>(ctx, a, b, out, "bitwise_or.Tensor_out");
}

Tensor& bitwise_or_Scalar_out(
Expand Down Expand Up @@ -141,7 +77,8 @@ Tensor& bitwise_or_Scalar_out(
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value = bitwise_or(a_casted, b_casted);
CTYPE_IN value =
std::bit_or<CTYPE_IN>()(a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
Expand Down
78 changes: 6 additions & 72 deletions kernels/portable/cpu/op_bitwise_xor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,10 @@
* LICENSE file in the root directory of this source tree.
*/

#include <cmath>
// patternlint-disable-next-line executorch-cpp-nostdinc
#include <functional>

#include <executorch/kernels/portable/cpu/pattern/bitwise_op.h>
#include <executorch/kernels/portable/cpu/scalar_utils.h>
#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
#include <executorch/kernels/portable/cpu/util/functional_util.h>
Expand All @@ -17,82 +19,14 @@ namespace torch {
namespace executor {
namespace native {

namespace {

template <typename CTYPE>
CTYPE bitwise_xor(CTYPE a, CTYPE b) {
return a ^ b;
}

template <>
bool bitwise_xor<bool>(bool a, bool b) {
return a != b;
}

} // namespace

using Tensor = exec_aten::Tensor;

Tensor& bitwise_xor_Tensor_out(
RuntimeContext& ctx,
const Tensor& a,
const Tensor& b,
Tensor& out) {
// Determine output size and resize for dynamic shapes
ET_KERNEL_CHECK(
ctx,
resize_to_broadcast_target_size(a, b, out) == Error::Ok,
InvalidArgument,
out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = b.scalar_type();
ScalarType common_type = promoteTypes(a_type, b_type);
ScalarType out_type = out.scalar_type();

ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);

ET_SWITCH_INT_TYPES_AND(
Bool, a_type, ctx, "bitwise_xor.Tensor_out", CTYPE_A, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool, b_type, ctx, "bitwise_xor.Tensor_out", CTYPE_B, [&]() {
ET_SWITCH_INT_TYPES_AND(
Bool,
common_type,
ctx,
"bitwise_xor.Tensor_out",
CTYPE_IN,
[&]() {
ET_SWITCH_REAL_TYPES_AND(
Bool,
out_type,
ctx,
"bitwise_xor.Tensor_out",
CTYPE_OUT,
[&]() {
apply_binary_elementwise_fn<
CTYPE_A,
CTYPE_B,
CTYPE_OUT>(
[](const CTYPE_A val_a, const CTYPE_B val_b) {
CTYPE_IN a_casted =
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value =
bitwise_xor(a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
a,
b,
out);
});
});
});
});

return out;
return bitwise_op_out<std::bit_xor>(ctx, a, b, out, "bitwise_xor.Tensor_out");
}

Tensor& bitwise_xor_Scalar_out(
Expand Down Expand Up @@ -143,8 +77,8 @@ Tensor& bitwise_xor_Scalar_out(
static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted =
static_cast<CTYPE_IN>(val_b);
CTYPE_IN value =
bitwise_xor(a_casted, b_casted);
CTYPE_IN value = std::bit_xor<CTYPE_IN>()(
a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
Expand Down
62 changes: 62 additions & 0 deletions kernels/portable/cpu/pattern/bitwise_op.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
* All rights reserved.
*
* This source code is licensed under the BSD-style license found in the
* LICENSE file in the root directory of this source tree.
*/

#pragma once

#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
#include <executorch/runtime/kernel/kernel_includes.h>

namespace torch {
namespace executor {
namespace native {
template <template <typename> typename OpFunc>
Tensor& bitwise_op_out(
RuntimeContext& ctx,
const Tensor& a,
const Tensor& b,
Tensor& out,
const char* op_name) {
ET_KERNEL_CHECK(
ctx,
resize_to_broadcast_target_size(a, b, out) == Error::Ok,
InvalidArgument,
out);

ScalarType a_type = a.scalar_type();
ScalarType b_type = b.scalar_type();
ScalarType common_type = promoteTypes(a_type, b_type);
ScalarType out_type = out.scalar_type();

ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);

ET_SWITCH_INT_TYPES_AND(Bool, a_type, ctx, op_name, CTYPE_A, [&]() {
ET_SWITCH_INT_TYPES_AND(Bool, b_type, ctx, op_name, CTYPE_B, [&]() {
ET_SWITCH_INT_TYPES_AND(Bool, common_type, ctx, op_name, CTYPE_IN, [&]() {
ET_SWITCH_REAL_TYPES_AND(
Bool, out_type, ctx, op_name, CTYPE_OUT, [&]() {
apply_binary_elementwise_fn<CTYPE_A, CTYPE_B, CTYPE_OUT>(
[](const CTYPE_A val_a, const CTYPE_B val_b) {
CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
CTYPE_IN value = OpFunc<CTYPE_IN>()(a_casted, b_casted);

return static_cast<CTYPE_OUT>(value);
},
a,
b,
out);
});
});
});
});

return out;
}
} // namespace native
} // namespace executor
} // namespace torch
12 changes: 12 additions & 0 deletions kernels/portable/cpu/pattern/targets.bzl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,18 @@ def define_common_targets():
The directory containing this targets.bzl file should also contain both
TARGETS and BUCK files that call this function.
"""
runtime.cxx_library(
name = "bitwise_op",
exported_headers = [
"bitwise_op.h",
],
compiler_flags = [],
deps = [
"//executorch/runtime/kernel:kernel_includes",
],
visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/optimized/cpu/..."],
)

runtime.cxx_library(
name = "comparison_op",
exported_headers = [
Expand Down

0 comments on commit 1693060

Please sign in to comment.