From 35295ea7b03ba7636b382844bf6fe4ccb35757a4 Mon Sep 17 00:00:00 2001 From: Renat Idrisov Date: Sun, 22 Dec 2024 19:10:19 +0000 Subject: [PATCH 1/3] Fix inconsistent IR when printing R.nn.pad operator --- include/tvm/relax/attrs/nn.h | 1 + python/tvm/relax/transform/legalize_ops/nn.py | 2 +- src/relax/op/nn/nn.cc | 6 +++--- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h index e26cee26584b..0685e25dfbe0 100644 --- a/include/tvm/relax/attrs/nn.h +++ b/include/tvm/relax/attrs/nn.h @@ -562,6 +562,7 @@ struct AttentionAttrs : public tvm::AttrsNode { /*! \brief Attributes used for the padding operator */ struct PadAttrs : public tvm::AttrsNode { Array pad_width; + Expr pad_value; tvm::String pad_mode; TVM_DECLARE_ATTRS(PadAttrs, "relay.attrs.PadAttrs") { diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py index 8317d4504e1e..e17878fd098f 100644 --- a/python/tvm/relax/transform/legalize_ops/nn.py +++ b/python/tvm/relax/transform/legalize_ops/nn.py @@ -228,7 +228,7 @@ def _nn_pad(bb: BlockBuilder, call: Call) -> Expr: pad_after = pad_widths[1::2] return bb.call_te( topi.nn.pad, - call.args[0], + call.attrs.pad_value, pad_before=pad_before, pad_after=pad_after, pad_value=float(call.args[1].data.numpy()), diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index 7eccf47e4b06..f65c06103691 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -140,8 +140,9 @@ Expr pad(Expr data, Array pad_width, Expr pad_value, String pad_mode) { auto attrs = make_object(); attrs->pad_width = std::move(pad_width); attrs->pad_mode = std::move(pad_mode); + attrs->pad_value = std::move(pad_value); static const Op& op = Op::Get("relax.nn.pad"); - return Call(op, {data, pad_value}, Attrs(attrs), {}); + return Call(op, {data}, Attrs(attrs), {}); } TVM_REGISTER_GLOBAL("relax.op.nn.pad").set_body_typed(pad); @@ -171,9 +172,8 @@ StructInfo InferStructInfoPad(const Call& call, const BlockBuilder& ctx) { } TVM_REGISTER_OP("relax.nn.pad") - .set_num_inputs(2) + .set_num_inputs(1) .add_argument("data", "Tensor", "The input tensor.") - .add_argument("pad_value", "Tensor", "The value to fill in padded area with.") .set_attrs_type() .set_attr("FInferStructInfo", InferStructInfoPad) .set_attr("FPurity", Bool(true)); From 1677ec289d9665a82daf8440e13e16bd0a5b8454 Mon Sep 17 00:00:00 2001 From: Renat Idrisov Date: Sun, 26 Jan 2025 00:03:01 +0000 Subject: [PATCH 2/3] Fixing the build --- 3rdparty/flashinfer | 2 +- include/tvm/relax/attrs/nn.h | 4 +++- python/tvm/relax/transform/legalize_ops/nn.py | 4 ++-- src/relax/op/nn/nn.cc | 4 ++-- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/3rdparty/flashinfer b/3rdparty/flashinfer index a76ceedb9495..1e379898a589 160000 --- a/3rdparty/flashinfer +++ b/3rdparty/flashinfer @@ -1 +1 @@ -Subproject commit a76ceedb9495d3d05648c29a8e6bb45baa265f6c +Subproject commit 1e379898a589cdd4ff18a4621fcbe18d63501545 diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h index 0685e25dfbe0..cce5c38e1a5b 100644 --- a/include/tvm/relax/attrs/nn.h +++ b/include/tvm/relax/attrs/nn.h @@ -562,13 +562,15 @@ struct AttentionAttrs : public tvm::AttrsNode { /*! \brief Attributes used for the padding operator */ struct PadAttrs : public tvm::AttrsNode { Array pad_width; - Expr pad_value; + Optional pad_value; tvm::String pad_mode; TVM_DECLARE_ATTRS(PadAttrs, "relay.attrs.PadAttrs") { TVM_ATTR_FIELD(pad_width).describe( "Number of values padded to the edges of each axis, " "in the format of (before_1, after_1, ..., before_N, after_N)"); + TVM_ATTR_FIELD(pad_value).describe( + "The value to fill in padded area with"); TVM_ATTR_FIELD(pad_mode) .set_default("constant") .describe( diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py index e17878fd098f..69112fa478c6 100644 --- a/python/tvm/relax/transform/legalize_ops/nn.py +++ b/python/tvm/relax/transform/legalize_ops/nn.py @@ -228,10 +228,10 @@ def _nn_pad(bb: BlockBuilder, call: Call) -> Expr: pad_after = pad_widths[1::2] return bb.call_te( topi.nn.pad, - call.attrs.pad_value, + call.args[0], pad_before=pad_before, pad_after=pad_after, - pad_value=float(call.args[1].data.numpy()), + pad_value=float(call.attrs.pad_value.data.numpy()), primfunc_name_hint="pad", ) diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc index f65c06103691..9e2199e3b2be 100644 --- a/src/relax/op/nn/nn.cc +++ b/src/relax/op/nn/nn.cc @@ -136,11 +136,11 @@ TVM_REGISTER_OP("relax.nn.log_softmax") /* relax.nn.pad */ TVM_REGISTER_NODE_TYPE(PadAttrs); -Expr pad(Expr data, Array pad_width, Expr pad_value, String pad_mode) { +Expr pad(Expr data, Array pad_width, Optional pad_value, String pad_mode) { auto attrs = make_object(); attrs->pad_width = std::move(pad_width); - attrs->pad_mode = std::move(pad_mode); attrs->pad_value = std::move(pad_value); + attrs->pad_mode = std::move(pad_mode); static const Op& op = Op::Get("relax.nn.pad"); return Call(op, {data}, Attrs(attrs), {}); } From 9d45cdb4df3e24d43eef221ff0f1ed705a7718ef Mon Sep 17 00:00:00 2001 From: Renat Idrisov Date: Sun, 26 Jan 2025 03:12:38 +0000 Subject: [PATCH 3/3] Formatting --- include/tvm/relax/attrs/nn.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h index cce5c38e1a5b..36ba7d7a56a8 100644 --- a/include/tvm/relax/attrs/nn.h +++ b/include/tvm/relax/attrs/nn.h @@ -569,8 +569,7 @@ struct PadAttrs : public tvm::AttrsNode { TVM_ATTR_FIELD(pad_width).describe( "Number of values padded to the edges of each axis, " "in the format of (before_1, after_1, ..., before_N, after_N)"); - TVM_ATTR_FIELD(pad_value).describe( - "The value to fill in padded area with"); + TVM_ATTR_FIELD(pad_value).describe("The value to fill in padded area with"); TVM_ATTR_FIELD(pad_mode) .set_default("constant") .describe(