We shouldn't be changing the default as part of this PR i don't think

Oh sorry it was a draft PR. Just updated it, and now it's ready for review!

In this PR, aten::stack is not enabled for fusion yet. Will enable the fusion (w/wo dynamic shapes) in separate PRs.

Nit: let's be consistent in handling error cases. Two lines below we use TORCH_INTERNAL_ASSERT and here we're throwing an error.

Why do we generate both ifThenElse and CompareSelect?

This doesn't look right. First line makes

newAxes[:] = axes[:]

Second line then appends it with axes[dim+1:].

That means that if we have 10-d tensors and dim=0 the stacked tensor will have rank 19. I think it should be 11-d instead, am I missing something?

Also, please update the comment to describe what new axis are created by this op and where they are inserted.

The first line makes

newAxes = axes[:dim]

Overall, there will be N-1 elements in newAxes where N is the size of axes.

Ah, got it, it makes sense. Maybe it'd be easier to grasp if we call newAxes as outputAxes and axes as inputAxes - currently "new" corresponds to inputs which is somewhat counterintuitive (even though I understand why it's done that way - we're deducing input indexes from output indexes).

The test only checks stacking two tensors with equal dimensions. We should also check other cases, e.g.:

1. stacking just one tensor
2. stacking N>2 tensors
3. stacking tensors with different dimensions

It would be good to also have such tests in python, so that whenever opinfo covers stack we can use that too.

It might be also a good idea to have a test verifying IR that we generate.

The test only checks stacking two tensors with equal dimensions. We should also check other cases...

That's a good point! Will add more tests.

It would be good to also have such tests in python, so that whenever opinfo covers stack we can use that too.

Not sure if I understand this. Are you suggesting to test above cases from Python instead of C++? We don't have many unit tests for ops right now. I only know aten::sum, aten::cat, and aten::conv. For these 3 ops, we put the unit tests in 3 files under test/cpp/tensorexpr. Maybe we should put them together in one file like test_ops.cpp.

It might be also a good idea to have a test verifying IR that we generate.

Agree. Will add it.

Are you suggesting to test above cases from Python instead of C++?

Having C++ tests is fine, but currently the most comprehensive coverage that we have is in python:

It leverages OpInfo, which is basically a collection of representative inputs for an op. What I suggested was to add a test for stack to those tests as well (if possible). If stack is not supported by OpInfo, then I think it still would make sense to write a python test manually, as most of our ops are currently tested from python rather than from C++.

Thanks for the pointer!! Added such a test for aten::stack in PR #74077.

Can we add an implementation that lowers into multiple loops like in the case of aten::cat, as done in computeCatWoConditionals?

The conditionals based lowering does not give good perf on CPUs, mainly because we can't vectorize it. I think the conditionals based lowering still makes sense for CUDA.

The conditionals based lowering does not give good perf on CPUs, mainly because we can't vectorize it.

aten::stack only uses compareSelect conditionals which can be vectorized. I'll see if the other implementation has better performance. Will add it if it is.

Please use some representative shapes from the models you have been working on, for perf analysis.

I still believe the multiple loops approach would be faster than conditionals because we will have both spacial and temporal cache locality for all inputs. It would be good to confirm with a perf analysis.

If you want to do this as a followup, thats fine too.

Sounds great! I'll first enable aten::stack in the fusion, and then do a perf analysis on real models for both approaches.

It looks like function errors out when dim_ == axes.size(). Isn't that a correct input for stack?

The axes are for the output of aten::stack. For example,

a = [1, 2]
b = [3, 4]
output = aten::stack([a, b], dim=dim_)

axes.size() equals to 3, and dim_ should be a number between 0 to 2, or -1 to -3, right? It cannot be 3. We have such dim tests for stack.

Aah, I missed the point that axes is for the output of stack. Makes sense.

This check essentially just verifies that the output is 4-d. Do we want to check the actual IR?

Attempted to check the stmt IR in the following way,

CHECK-NEXT: aten_stack[Ramp(5ll * (j + 2ll * i), 1ll, 4)] = (Broadcast(j, 4))==(Broadcast(1ll, 4)) ? (ty_1[Ramp(5ll * i, 1ll, 4)]) : (tx_1[Ramp(5ll * i, 1ll, 4)]))

it failed one CI test (linux-xenial-py3.7-clang7-asan). With the configuration in this CI test, the vectorization failed thus the generated stmt is different. There are two ways to fix this,

1. implement a CHECK that can check patterns in a string such as ".?.:" for stack
2. save the original stmt before prepare_to_codegen in TEK for checking issues

Option 2 does not seem good as it introduces additional member vars in TEK only for testing issues; Option 1 can be implemented in a separate PR. Currently landing this PR with the original checks.