Skip to content

Commit

Permalink
pnnx convert torch bitwise left_shift right_shift (#4349)
Browse files Browse the repository at this point in the history
  • Loading branch information
nihui authored Nov 14, 2022
1 parent eceac35 commit 6967baa
Show file tree
Hide file tree
Showing 9 changed files with 258 additions and 22 deletions.
2 changes: 2 additions & 0 deletions tools/pnnx/src/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ set(pnnx_pass_level2_SRCS
pass_level2/torch_bitwise_and.cpp
pass_level2/torch_bitwise_or.cpp
pass_level2/torch_bitwise_xor.cpp
pass_level2/torch_bitwise_left_shift.cpp
pass_level2/torch_bitwise_right_shift.cpp
pass_level2/torch_cat.cpp
pass_level2/torch_chunk.cpp
pass_level2/torch_clamp.cpp
Expand Down
4 changes: 3 additions & 1 deletion tools/pnnx/src/ir.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1134,7 +1134,7 @@ static std::string expand_expression(const Operator* op)
std::string r = binaryop + "(" + a + ", " + b + ")";
exprstack.push(r);
}
else if (t == "add" || t == "sub" || t == "mul" || t == "div" || t == "floor_divide" || t == "and" || t == "or" || t == "xor")
else if (t == "add" || t == "sub" || t == "mul" || t == "div" || t == "floor_divide" || t == "and" || t == "or" || t == "xor" || t == "lshift" || t == "rshift")
{
std::string binaryop;
if (t == "add") binaryop = "+";
Expand All @@ -1145,6 +1145,8 @@ static std::string expand_expression(const Operator* op)
if (t == "and") binaryop = "&";
if (t == "or") binaryop = "|";
if (t == "xor") binaryop = "^";
if (t == "lshift") binaryop = "<<";
if (t == "rshift") binaryop = ">>";

std::string a = exprstack.top();
exprstack.pop();
Expand Down
41 changes: 41 additions & 0 deletions tools/pnnx/src/pass_level2/torch_bitwise_left_shift.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include "pass_level2.h"

namespace pnnx {

class torch_bitwise_left_shift : public GraphRewriterPass
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
4 3
pnnx.Input input_0 0 1 input
pnnx.Input input_1 0 1 other
aten::bitwise_left_shift op_0 2 1 input other out
pnnx.Output output 1 0 out
)PNNXIR";
}

const char* type_str() const
{
return "torch.bitwise_left_shift";
}
};

REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_bitwise_left_shift, 20)

} // namespace pnnx
41 changes: 41 additions & 0 deletions tools/pnnx/src/pass_level2/torch_bitwise_right_shift.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#include "pass_level2.h"

namespace pnnx {

class torch_bitwise_right_shift : public GraphRewriterPass
{
public:
const char* match_pattern_graph() const
{
return R"PNNXIR(7767517
4 3
pnnx.Input input_0 0 1 input
pnnx.Input input_1 0 1 other
aten::bitwise_right_shift op_0 2 1 input other out
pnnx.Output output 1 0 out
)PNNXIR";
}

const char* type_str() const
{
return "torch.bitwise_right_shift";
}
};

REGISTER_GLOBAL_PNNX_GRAPH_REWRITER_PASS(torch_bitwise_right_shift, 20)

} // namespace pnnx
10 changes: 4 additions & 6 deletions tools/pnnx/src/pass_level3/fuse_expression.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ static bool operand_maybe_tensor(const Operand* operand)
return operand_maybe_tensor(op->inputs[0]) || operand_maybe_tensor(op->inputs[1]);
}

if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__")
if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__" || op->type == "aten::__lshift__" || op->type == "aten::__rshift__")
{
return operand_maybe_tensor(op->inputs[0]) || operand_maybe_tensor(op->inputs[1]);
}
Expand Down Expand Up @@ -321,11 +321,9 @@ static void fuse_expression(Graph& graph, Operand* operand, std::string& expr, s
fuse_expression(graph, op->inputs[1], expr, inputs, foldable_constants);
expr += ")";
}
else if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__")
else if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__" || op->type == "aten::__lshift__" || op->type == "aten::__rshift__")
{
std::string mathop = op->type.substr(8, 3);
if (mathop == "or_")
mathop = "or";
std::string mathop = op->type.substr(8, op->type.size() - 10);

expr += mathop;
expr += "(";
Expand Down Expand Up @@ -485,7 +483,7 @@ void fuse_expression(Graph& graph, const std::set<std::string>& foldable_constan
{
need_fuse = true;
}
if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__")
if (op->type == "aten::__and__" || op->type == "aten::__or__" || op->type == "aten::__xor__" || op->type == "aten::__lshift__" || op->type == "aten::__rshift__")
{
need_fuse = true;
}
Expand Down
67 changes: 52 additions & 15 deletions tools/pnnx/src/pass_level5/eval_expression.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -45,20 +45,14 @@ static bool token_is_literal(const std::string& t)
float f;
iss >> std::noskipws >> f;
return iss.eof() && !iss.fail();
}

// for (size_t i = 0; i < t.size(); i++)
// {
// if (i == 0 && t[i] == '-')
// continue;
//
// if (t[i] < '0' || t[i] > '9')
// {
// if (t[i] != '.' && t[i] != 'e')
// return false;
// }
// }
//
// return true;
static bool token_is_interger_literal(const std::string& t)
{
std::istringstream iss(t);
int f;
iss >> std::noskipws >> f;
return iss.eof() && !iss.fail();
}

static std::string eval_expression(const Operator* op)
Expand Down Expand Up @@ -317,8 +311,7 @@ static std::string eval_expression(const Operator* op)
|| t == "div"
|| t == "floor_divide"
|| t == "pow"
|| t == "remainder"
|| t == "and" || t == "or" || t == "xor")
|| t == "remainder")
{
std::string a = exprstack.top();
exprstack.pop();
Expand Down Expand Up @@ -379,6 +372,50 @@ static std::string eval_expression(const Operator* op)
exprstack.push(r);
}
}
else if (t == "and" || t == "or" || t == "xor" || t == "lshift" || t == "rshift")
{
std::string a = exprstack.top();
exprstack.pop();
std::string b = exprstack.top();
exprstack.pop();

if (token_is_interger_literal(a) && token_is_interger_literal(b))
{
int ai = std::stoi(a);
int bi = std::stoi(b);

if (t == "and")
{
int r = ai & bi;
exprstack.push(std::to_string(r));
}
if (t == "or")
{
int r = ai | bi;
exprstack.push(std::to_string(r));
}
if (t == "xor")
{
int r = ai ^ bi;
exprstack.push(std::to_string(r));
}
if (t == "lshift")
{
int r = ai << bi;
exprstack.push(std::to_string(r));
}
if (t == "rshift")
{
int r = ai >> bi;
exprstack.push(std::to_string(r));
}
}
else
{
std::string r = t + "(" + a + "," + b + ")";
exprstack.push(r);
}
}
else if (t == "[") // list
{
std::vector<std::string> elements;
Expand Down
5 changes: 5 additions & 0 deletions tools/pnnx/tests/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,11 @@ if(Torch_VERSION VERSION_GREATER_EQUAL "1.9")
pnnx_add_test(nn_Mish)
endif()

if(Torch_VERSION VERSION_GREATER_EQUAL "1.10")
pnnx_add_test(torch_bitwise_left_shift)
pnnx_add_test(torch_bitwise_right_shift)
endif()

if(Torch_VERSION VERSION_GREATER_EQUAL "1.11")
pnnx_add_test(torch_fft_ihfft2)
pnnx_add_test(torch_fft_ihfftn)
Expand Down
55 changes: 55 additions & 0 deletions tools/pnnx/tests/test_torch_bitwise_left_shift.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# Tencent is pleased to support the open source community by making ncnn available.
#
# Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()

def forward(self, x, y):
out = torch.bitwise_left_shift(x, y)
return out

def test():
net = Model()
net.eval()

torch.manual_seed(0)
x = torch.randint(10, (3, 16), dtype=torch.int)
y = torch.randint(10, (3, 16), dtype=torch.int)

a = net(x, y)

# export torchscript
mod = torch.jit.trace(net, (x, y))
mod.save("test_torch_bitwise_left_shift.pt")

# torchscript to pnnx
import os
os.system("../src/pnnx test_torch_bitwise_left_shift.pt inputshape=[3,16]i32,[3,16]i32")

# pnnx inference
import test_torch_bitwise_left_shift_pnnx
b = test_torch_bitwise_left_shift_pnnx.test_inference()

return torch.equal(a, b)

if __name__ == "__main__":
if test():
exit(0)
else:
exit(1)
55 changes: 55 additions & 0 deletions tools/pnnx/tests/test_torch_bitwise_right_shift.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# Tencent is pleased to support the open source community by making ncnn available.
#
# Copyright (C) 2022 THL A29 Limited, a Tencent company. All rights reserved.
#
# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
#
# https://opensource.org/licenses/BSD-3-Clause
#
# Unless required by applicable law or agreed to in writing, software distributed
# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
# CONDITIONS OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the License.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()

def forward(self, x, y):
out = torch.bitwise_right_shift(x, y)
return out

def test():
net = Model()
net.eval()

torch.manual_seed(0)
x = torch.randint(10, (3, 16), dtype=torch.int)
y = torch.randint(10, (3, 16), dtype=torch.int)

a = net(x, y)

# export torchscript
mod = torch.jit.trace(net, (x, y))
mod.save("test_torch_bitwise_right_shift.pt")

# torchscript to pnnx
import os
os.system("../src/pnnx test_torch_bitwise_right_shift.pt inputshape=[3,16]i32,[3,16]i32")

# pnnx inference
import test_torch_bitwise_right_shift_pnnx
b = test_torch_bitwise_right_shift_pnnx.test_inference()

return torch.equal(a, b)

if __name__ == "__main__":
if test():
exit(0)
else:
exit(1)

0 comments on commit 6967baa

Please sign in to comment.