Question on check_transpose

[zcbenz]

When checking whether a matrix is transposed the implementation in meta/matmul.cpp has a check for !is_vector:

mlx/mlx/backend/metal/matmul.cpp

Lines 93 to 94 in 9eb7d73

	if (sty == 1 && (!is_vector || stx == arr.shape(-1))) {
	return std::make_tuple(false, stx, arr);

Shouldn't sty == 1 && stx == arr.shape(-1) be enough? Or is there an edge case to be handled?

The cpu/matmul.cpp does so:

mlx/mlx/backend/cpu/matmul.cpp

Lines 20 to 21 in 9eb7d73

	if (stx == arr.shape(-1) && sty == 1) {
	return std::make_tuple(false, stx, arr);

[awni]

I think the function is poorly named for what it does. It really does a check + optional copy in certain cases. And those cases are different for the CPU and GPU since the GPU can support irregular strides for the last two dimensions (in some cases).

• For non-vectors checks if it's transposed. We support irregular strides in some cases. So the condition stx == arr.shape(-1) (for example) does not need to be satisfied for non-vectors
• For vectors I assume we don't support irregular strides in the vector dimension so it will fallback to do a copy to make it contiguous

For the CPU there is no irregular stride support in either of the last two dimensions so we always do a copy if the basic conditions are not met.

[awni]

Here's a test case you can use:

import mlx.core as mx

a = mx.broadcast_to(mx.array([1.0, 2.0]), (2, 2))
b = mx.array([3.0, 4.0])
print(a @ b)

[zcbenz]

broadcast_to still does a copy, you can add a fprintf(stderr, "a: %zu, b: %zu\n", a.size(), b.size()); in matmul to see it.

[awni]

size is not the true size of the data. Use data_size for that.

[awni]

Broadcast never copies btw.

[zcbenz]

Ah I get it now: when advancing the matrix, gemv kernel uses out_row * matrix_ld, and matrix_ld is actually the stride instead of dimension, which works with irregular strides.

Thanks for your help!