Lower as_strided_copy use fast path with slice

[zpcore]

When we execute the following two code snippets regarding flash attention kernel in custom_kernel.py, they suppose to produce the same result.
a.

l, m = (v[..., 0] for v in aux[-2:]) 

b.

l, m = aux[-2:] l = torch.ops.aten.slice(l, -1, 0, 1) m = torch.ops.aten.slice(m, -1, 0, 1) 

Both will be lowered through

xla/torch_xla/csrc/aten_xla_type.cpp

Line 882 in 1acc987

at::Tensor XLANativeFunctions::as_strided_copy(

, the difference is that input argument stride and size will be one element fewer in code a compared with code b. With such argument difference, code a will be fallback into aten::take and this can trigger the following error when we call with SPMD:

F0223 07:18:45.157172 842998 hlo_sharding.cc:1024] Check failed: !IsManual() 

I plan to check in test_as_stride_use_slice.py in this PR.
Note 1. Failing test test_scan_layer_aot is not enabled until #8742 is resolved.
2. Failing test test_scan_weight_layer_aot is not enabled unti #8753 is resolved

[tengyifei]

Should we check that stride[j] can be evenly divided by stride_mul and exit if the remainder is not 0?

[zpcore]

It doesn't need to be 'evenly divided' for stride[j] as long as all indexes before j of stride matches with the cumulative product of tensor dim.