Skip to content

[Prim] Add p_norm_decomp #63156

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
HydrogenSulfate merged 10 commits into from
Apr 7, 2024
Merged

[Prim] Add p_norm_decomp #63156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
95f6312
add p_norm_decomp
HydrogenSulfate Apr 1, 2024
bd1f3e3
use rsqrt
HydrogenSulfate Apr 1, 2024
4c7c702
support for porder is float
HydrogenSulfate Apr 1, 2024
d8fb386
update UT
HydrogenSulfate Apr 2, 2024
33e6a23
update UT
HydrogenSulfate Apr 2, 2024
32243a4
update atol and rtol
HydrogenSulfate Apr 2, 2024
363261a
support 0, 1, inf, -inf norm
HydrogenSulfate Apr 2, 2024
a097848
update UT for 0, 1, inf, -inf norm
HydrogenSulfate Apr 2, 2024
cc062a0
Merge branch 'PaddlePaddle:develop' into add_p_norm_decomp
HydrogenSulfate Apr 5, 2024
6f77683
remove 0-norm for it has benn replaced with paddle.count_nonzero in l…
HydrogenSulfate Apr 5, 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
2 changes: 2 additions & 0 deletions paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
"leaky_relu",
"log_softmax",
"mean",
"p_norm",
"pow",
"relu",
"relu6",
Expand Down Expand Up @@ -72,6 +73,7 @@
"leaky_relu",
"log_softmax",
"mean",
"p_norm",
"pow",
"relu",
"relu6",
Expand Down
54 changes: 54 additions & 0 deletions paddle/fluid/primitive/composite/composite.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,60 @@ static bool valid_type(const DataType& dtype) {
}
}

template <typename T>
Tensor p_norm_decomp(const Tensor& x,
const float& porder = 2.0,
const int& axis = -1,
const float epsilon = 1.0e-12f,
const bool& keepdim = false,
const bool& asvector = false) {
auto org_dtype = x.dtype();
auto x_tmp = x;

bool need_cast = is_half_dtype(org_dtype);
if (need_cast) {
x_tmp = cast<T>(x, DataType::FLOAT32);
}

Tensor res;
if (porder == 0.0) {
// 0-norm
auto zero = full<T>(empty_shape, 0, x_tmp.dtype());
auto none_zero = not_equal<T>(x_tmp, zero);
res = cast<T>(none_zero, x_tmp.dtype());
res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
} else if (porder == 1.0) {
// 1-norm
res = abs<T>(x_tmp);
res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
} else if (porder == 2.0) {
// 2-norm
auto one = full<T>(empty_shape, 1, x_tmp.dtype());
res = one / rsqrt<T>(sum<T>(x_tmp * x_tmp, {axis}, x_tmp.dtype(), keepdim));
} else if (porder == INFINITY) {
// +INF-norm
res = abs<T>(x_tmp);
res = max<T>(x_tmp, {axis}, keepdim);
} else if (porder == -INFINITY) {
// -INF-norm
res = abs<T>(x_tmp);
res = min<T>(x_tmp, {axis}, keepdim);
} else {
// vanilla p-norm
auto porder_tensor = full<T>(empty_shape, porder, x_tmp.dtype());
auto inv_porder_tensor = full<T>(empty_shape, 1 / porder, x_tmp.dtype());
res = elementwise_pow<T>(x_tmp, porder_tensor);
res = sum<T>(res, {axis}, x_tmp.dtype(), keepdim);
res = elementwise_pow<T>(res, inv_porder_tensor);
}

if (need_cast) {
return cast<T>(res, org_dtype);
} else {
return res;
}
}

Copy link
Contributor

@cxxly cxxly Apr 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

建议组合算子签名和PHI/PIR签名保持一致，用namespace约束分类，函数名不要加_decomp。存量的可以后续统一修改

template <typename T>
Tensor pow_decomp(const Tensor& x, const paddle::Scalar& y) {
auto org_dtype = x.dtype();
Expand Down
40 changes: 33 additions & 7 deletions test/legacy_test/test_norm_all.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,17 @@ def p_norm_python_api(
return _C_ops.p_norm(x, p, axis, epsilon, keepdim, as_vector)


def norm_public_python_api(
x, p=2.0, axis=-1, epsilon=1e-12, keepdim=False, as_vector=False
):
return paddle.linalg.norm(
x,
p,
axis,
keepdim,
)


def np_linalg_vector_norm(x, axis, porder, keepdims=False):
x_shape = list(x.shape)

Expand Down Expand Up @@ -175,8 +186,14 @@ class TestPnormOp(OpTest):
def setUp(self):
self.op_type = "p_norm"
self.python_api = p_norm_python_api
self.public_python_api = norm_public_python_api
self.prim_op_type = "comp"
self.init_test_case()
self.init_dtype()
self.fw_comp_atol = 1e-6
self.fw_comp_rtol = 1e-6
self.rev_comp_atol = 1e-6
self.rev_comp_rtol = 1e-6
x = (np.random.random(self.shape) + 0.5).astype(self.dtype)
norm = np_linalg_norm(x, self.axis, self.porder, self.keepdim)
self.inputs = {'X': x}
Expand All @@ -191,10 +208,10 @@ def setUp(self):
self.gradient = self.calc_gradient()

def test_check_output(self):
self.check_output()
self.check_output(check_prim_pir=True)

def test_check_grad(self):
self.check_grad(['X'], 'Out')
self.check_grad(['X'], 'Out', check_prim_pir=True)

def init_test_case(self):
self.shape = [2, 3, 4, 5]
Expand Down Expand Up @@ -257,7 +274,7 @@ def init_dtype(self):
self.dtype = "float32"

def test_check_grad(self):
self.check_grad(['X'], 'Out')
self.check_grad(['X'], 'Out', check_prim_pir=True)


class TestPnormOp3(TestPnormOp):
Expand All @@ -273,7 +290,9 @@ def init_dtype(self):
self.dtype = "float32"

def test_check_grad(self):
self.check_grad(['X'], 'Out', user_defined_grads=self.gradient)
self.check_grad(
['X'], 'Out', user_defined_grads=self.gradient, check_prim_pir=True
)


class TestPnormOp4(TestPnormOp):
Expand All @@ -289,7 +308,9 @@ def init_dtype(self):
self.dtype = "float32"

def test_check_grad(self):
self.check_grad(['X'], 'Out', user_defined_grads=self.gradient)
self.check_grad(
['X'], 'Out', user_defined_grads=self.gradient, check_prim_pir=True
)


class TestPnormOp5(TestPnormOp):
Expand Down Expand Up @@ -321,7 +342,9 @@ def init_dtype(self):
self.dtype = "float32"

def test_check_grad(self):
self.check_grad(['X'], 'Out', user_defined_grads=self.gradient)
self.check_grad(
['X'], 'Out', user_defined_grads=self.gradient, check_prim_pir=True
)


def create_test_fp16_class(parent, max_relative_error=2e-3):
Expand Down Expand Up @@ -367,7 +390,9 @@ def test_check_grad(self):
class TestPnormBF16Op(OpTest):
def setUp(self):
self.op_type = "p_norm"
self.prim_op_type = "comp"
self.python_api = p_norm_python_api
self.public_python_api = norm_public_python_api
self.init_test_case()
self.x = (np.random.random(self.shape) + 0.5).astype(np.float32)
self.norm = np_linalg_norm(self.x, self.axis, self.porder, self.keepdim)
Expand All @@ -384,7 +409,7 @@ def setUp(self):

def test_check_output(self):
place = core.CUDAPlace(0)
self.check_output_with_place(place, atol=1e-3)
self.check_output_with_place(place, atol=1e-3, check_prim_pir=True)

def test_check_grad(self):
place = core.CUDAPlace(0)
Expand All @@ -393,6 +418,7 @@ def test_check_grad(self):
['X'],
'Out',
user_defined_grads=self.gradient,
check_prim_pir=True,
)

def init_test_case(self):
Expand Down