[Relay, TOPI] Refactor strided_slice and add axes argument

[masahi]

This PR has two purposes:

1. Add axes argument to static param strided_slice
2. Refactor strided_slice related code so that adding axes argument would not increase clutter. Also I did further refactoring to remove duplicated output shape calculation between topi/relay type rel (left as TODO) and minimize the number of te::compute for strided slice (currently there are 4 of them).

For 1, currently slicing along the second axis with input shape (?, 3), begin [0], end [1] calls the dynamic parameter variant of relay/topi strided_slice with begin = (0, 0) and end = (?, ?), which results in the output shape (?, ?) rather than desired (?, 1). This causes a problem when the output is fed to gather_nd which requires a certain dimension to be static.

The root problem is that our strided_slice op does not have a notion of axes and always tries to slice along all axes. So when some of the dimensions are dynamic, we call dynamic parameter variant of strided_slice which makes all input params and the output shape dynamic, even if we only need to slice along static dimensions with static begin/end. My solution adds axes argument to static parameter variant of strided_sliceso that slicing would not touch dynamic dimensions and preserve static dimensions along provided axes.

For 2, we accumulated many patches for strided_slice over the last few years and I found the current implementation extremely messy and hard to maintain. In order not to make the situation worse by introducing axes argument, I decided to roll major refactoring of our strided_slice related code.

First, the output shape calculation logic is duplicated between topi and relay type rel:

• https://github.com/apache/tvm/blob/main/include/tvm/topi/transform.h#L651-L714
• https://github.com/apache/tvm/blob/main/src/relay/op/tensor/transform.cc#L2451-L2534

This has been fixed by introducing StridedSliceOutputShape function in topi and use it from both topi/relay.

Second, currently there are 4 te::compute for strided_slice, each with slightly different expectation on its input and input param canonicalization.

• https://github.com/apache/tvm/blob/main/include/tvm/topi/transform.h#L716-L726
• https://github.com/apache/tvm/blob/main/include/tvm/topi/transform.h#L628-L636
• https://github.com/apache/tvm/blob/main/include/tvm/topi/transform.h#L579-L593
• https://github.com/apache/tvm/blob/main/src/relay/op/tensor/transform.cc#L2710-L2720

I examined each of them in detail and decided that we only need two te::compute to support all cases:

• A static parameter variant, where all input parameters are Array<Integer>. Input shape can be partially dynamic/static.

  tvm/include/tvm/topi/transform.h

  Lines 700 to 712 in 64c1bbf

  	return te::compute(
  	out_shape,
  	[&](const Array<tir::Var>& indices) {
  	Array<PrimExpr> real_indices;
  	for (size_t i = 0; i < out_shape.size(); ++i) real_indices.push_back(indices[i]);
  	for (size_t i = 0; i < axes.size(); ++i) {
  	auto stride = make_const(strides[i].dtype(), strides_vec[i]);
  	PrimExpr ind = indices[axes[i]] * stride + begin_expr[i];
  	real_indices.Set(axes[i], ind);
  	}
  	return x(real_indices);
  	},
  	name, tag);

• A partially dynamic/static parameter variant, where all input parameters are Array<PrimExpr>.

  tvm/include/tvm/topi/transform.h

  Lines 594 to 608 in 64c1bbf

  	return te::compute(
  	out_shape,
  	[&](const Array<tvm::tir::Var>& indices) {
  	Array<PrimExpr> real_indices;
  	for (size_t i = 0; i < num_slice_axes; ++i) {
  	real_indices.push_back(indices[i] * strides[i] + tvm::min(begin[i], x->shape[i] - 1));
  	}
  	// keep input dim
  	for (size_t i = num_slice_axes; i < src_tensor_dim; ++i) {
  	real_indices.push_back(indices[i]);
  	}
  	return x(real_indices);
  	},
  	name, tag);
  	}

In both cases, static input dimensions are preserved as much as possible. The axes is added only to the first variant, according to my use case.

please review @tqchen @mbrookhart @kevinthesun @yongwww @comaniac

[mbrookhart]

Mostly LGTM, a couple of nitpicks. My only design complaint is that we seem to have some API consistency problems. If we provide a dynamic begin and an axes argument, it looks like we'll hit a codepath that doesn't use the axes properly?

[mbrookhart]

Do we not support axes with dynamic begin?

[mbrookhart]

Perhaps we should throw an error if we hit this codepath and axes is defined? With a TODO for extending axes to the dynamic case?

[masahi]

Yeah I haven't put much thought into axes argument for dynamic strided slice. My goal was to preserve more static dimensions along axes, which doesn't apply to the dynamic parameter variant. We should support it for convenience and API consistency sake, but for now I'd like to leave it as TODO and assert axes is None here until someone complains or comes up with a good use case.

[masahi]

Added TODO and assert.

[masahi]

thanks @mbrookhart