gemini-code-assist[bot] commented on code in PR #18622:
URL: https://github.com/apache/tvm/pull/18622#discussion_r2650209948
##########
src/relax/op/tensor/manipulate.cc:
##########
@@ -1929,12 +1929,21 @@ StructInfo InferStructInfoTile(const Call& call, const
BlockBuilder& ctx) {
return TensorStructInfo(ShapeExpr(out_shape), data_sinfo->dtype,
data_sinfo->vdevice);
}
-// TODO(relax-team): implement FRelaxInferLayout for tile
+InferLayoutOutput InferLayoutTile(
+ const Call& call, const ffi::Map<ffi::String, ffi::Array<ffi::String>>&
desired_layouts,
+ const VarLayoutMap& var_layout_map) {
+ ICHECK(NoDesiredLayout(call, desired_layouts));
+ LayoutDecision existing_layout = GetLayoutDecision(var_layout_map,
call->args[0]);
+
+ return InferLayoutOutput({existing_layout}, {existing_layout},
Attrs(call->attrs));
+}
Review Comment:
The current implementation of `InferLayoutTile` incorrectly assumes that the
`repeats` attribute of the `tile` operator does not need to be transformed when
the input layout changes. It simply propagates the input layout to the output
and leaves the attributes unchanged.
This is incorrect because the `repeats` array is positional and corresponds
to the dimensions of the input tensor. When the `ConvertLayout` pass permutes
the dimensions of the input tensor, the `repeats` array should also be permuted
to match the new dimension order to preserve the semantics.
For example, if we have an `NCHW` tensor and we want to tile the `C` axis
(axis 1), `repeats` would be `(1, 2, 1, 1)`. If the layout is changed to
`NHWC`, the `C` axis is now at index 3. To tile the same logical `C` axis,
`repeats` must be transformed to `(1, 1, 1, 2)`. The current implementation
fails to do this.
This can lead to silent incorrect transformations and wrong results.
I suggest a more complete implementation that handles `repeats`
transformation, including cases where dimensions are extended.
```c
InferLayoutOutput InferLayoutTile(
const Call& call, const ffi::Map<ffi::String, ffi::Array<ffi::String>>&
desired_layouts,
const VarLayoutMap& var_layout_map) {
ICHECK(NoDesiredLayout(call, desired_layouts));
const auto* attrs = call->attrs.as<TileAttrs>();
ICHECK(attrs);
LayoutDecision existing_layout = GetLayoutDecision(var_layout_map,
call->args[0]);
if (existing_layout.IsUnknown()) {
auto* in_sinfo = call->args[0]->struct_info_.as<TensorStructInfoNode>();
ICHECK(in_sinfo);
auto* out_sinfo = GetStructInfo(call).as<TensorStructInfoNode>();
ICHECK(out_sinfo);
return InferLayoutOutput({InitialLayoutDecision(in_sinfo->ndim)},
{InitialLayoutDecision(out_sinfo->ndim)},
Attrs(attrs));
}
// The logic to handle dimension extension from numpy.tile
int ndim = existing_layout->layout.ndim();
int repeat_size = attrs->repeats.size();
int out_ndim = std::max(ndim, repeat_size);
Layout initial_layout_extended = InitialLayout(out_ndim);
Layout existing_layout_extended = existing_layout->layout;
if (ndim < out_ndim) {
std::string layout_str = InitialLayout(out_ndim - ndim).name();
for (int i = 0; i < ndim; ++i) {
layout_str.push_back(existing_layout->layout.name()[i] + (out_ndim -
ndim));
}
existing_layout_extended = Layout(layout_str);
}
Array<Integer> repeats_extended;
for (int i = 0; i < out_ndim - repeat_size; ++i) {
repeats_extended.push_back(Integer(1));
}
for (const auto& val : attrs->repeats) {
repeats_extended.push_back(val);
}
Array<Integer> new_repeats;
for (int i = 0; i < out_ndim; ++i) {
new_repeats.push_back(
repeats_extended[initial_layout_extended.IndexOf(existing_layout_extended[i])]);
}
auto new_attrs = make_object<TileAttrs>();
new_attrs->repeats = new_repeats;
return InferLayoutOutput({existing_layout}, {existing_layout},
Attrs(new_attrs));
}
```
##########
tests/python/relax/test_transform_convert_layout.py:
##########
@@ -5077,5 +5077,49 @@ def main(
verify(Input, Expected)
+def test_conv2d_tile():
+ @I.ir_module
+ class Input:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3),
"float32")
+ ) -> R.Tensor(None, "float32", ndim=4):
+ with R.dataflow():
+ gv: R.Tensor((2, 4, 26, 26), "float32") = R.nn.conv2d(x, w,
out_dtype="float32")
+ gv2: R.Tensor((4, 4, 26, 26), "float32") = R.tile(gv,
repeats=(2, 1, 1, 1))
+ R.output(gv2)
+ return gv2
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3, 28, 28), dtype="float32"), w: R.Tensor((4, 3,
3, 3), dtype="float32")
+ ) -> R.Tensor(None, dtype="float32", ndim=4):
+ with R.dataflow():
+ lv: R.Tensor((2, 28, 28, 3), dtype="float32") =
R.permute_dims(x, axes=[0, 2, 3, 1])
+ lv1: R.Tensor((4, 3, 3, 3), dtype="float32") =
R.permute_dims(w, axes=[0, 2, 3, 1])
+ gv: R.Tensor((2, 26, 26, 4), dtype="float32") = R.nn.conv2d(
+ lv,
+ lv1,
+ strides=[1, 1],
+ padding=[0, 0, 0, 0],
+ dilation=[1, 1],
+ groups=1,
+ data_layout="NHWC",
+ kernel_layout="OHWI",
+ out_layout="NHWC",
+ out_dtype="float32",
+ )
+ lv2: R.Tensor((4, 26, 26, 4), dtype="float32") = R.tile(gv,
repeats=(2, 1, 1, 1))
+ gv2: R.Tensor((4, 4, 26, 26), dtype="float32") =
R.permute_dims(
+ lv2, axes=[0, 3, 1, 2]
+ )
+ R.output(gv2)
+ return gv2
+
+ verify(Input, Expected)
Review Comment:
The added test case `test_conv2d_tile` is good, but it doesn't fully cover
the correctness of the `InferLayoutTile` implementation. The test passes
because it tiles along axis 0 (the `N` dimension), which has the same index in
both `NCHW` and `NHWC` layouts.
A more robust test would be to tile along an axis that is moved during the
layout transformation, for example, the channel axis (`C`). This would expose
the bug in the current `InferLayoutTile` implementation which fails to
transform the `repeats` attribute.
Consider a test case where you tile on the channel axis:
- In the `Input` module, use `R.tile(gv, repeats=(1, 2, 1, 1))` to tile the
channel axis of the `NCHW` tensor.
- In the `Expected` module, the `repeats` for `R.tile` on the `NHWC` tensor
should be transformed to `(1, 1, 1, 2)` to correspond to the new position of
the channel axis.
This would ensure that the layout inference for `tile` correctly handles
attribute transformation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
