Skip to content

GitLab

"include/vscode:/vscode.git/clone" did not exist on "0381f2872cabc3e675bae16fcfa483c9961f46e9"
Commit 8a2837c4 authored by Alan Turner's avatar Alan Turner
Browse files

Add where op workaround for CK

parent a47adf4b
Hide whitespace changes
Inline Side-by-side
Showing with 9 additions and 0 deletions
src/onnx/parse_where.cpp
View file @ 8a2837c4
...@@ -31,6 +31,8 @@ namespace migraphx { ...@@ -31,6 +31,8 @@ namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS { inline namespace MIGRAPHX_INLINE_NS {
namespace onnx { namespace onnx {
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_ENABLE_CK)
struct parse_where : op_parser<parse_where> struct parse_where : op_parser<parse_where>
{ {
std::vector<op_desc> operators() const { return {{"Where"}}; } std::vector<op_desc> operators() const { return {{"Where"}}; }
...@@ -56,6 +58,13 @@ struct parse_where : op_parser<parse_where> ...@@ -56,6 +58,13 @@ struct parse_where : op_parser<parse_where>
auto lens = auto lens =
compute_broadcasted_lens(args[0]->get_shape().lens(), args[1]->get_shape().lens()); compute_broadcasted_lens(args[0]->get_shape().lens(), args[1]->get_shape().lens());
lens = compute_broadcasted_lens(lens, args[2]->get_shape().lens()); lens = compute_broadcasted_lens(lens, args[2]->get_shape().lens());
if (enabled(MIGRAPHX_ENABLE_CK{}))
{
// Convert condition tensor to int32 to work around CK not supporting bool type
args[0] = info.add_instruction(make_op("convert", {{"target_type", shape::int32_type}}), args[0]);
}
if(args[0]->get_shape().lens() != lens) if(args[0]->get_shape().lens() != lens)
{ {
args[0] = args[0] =
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment