[BUG]: Weird behavior when passing a tensor as owned to a function

[andresnowak]

Bug description

When passing a tensor as owned to a function and one tries to do a memcpy of the data or printing the information from inside a @parameter function (using a simd load) a weird behaviors happens where the first data values are trash values.

Steps to reproduce

Code

fn foo(owned tensor: Tensor[dtype]):
    print("first", tensor[0])
    print("second", tensor.load[width=4](0))
 
    @parameter
    fn bar[width: Int](i: Int):
        print("third", tensor.load[width=width](i))
 
    bar[width=4](0)
 
fn foo_2(owned tensor: Tensor[dtype]):
    print("first", tensor[0])
    print("second", tensor.load[width=4](0))
 
    @parameter
    fn bar[width: Int](i: Int):
        print("third", tensor.load[width=width](i))
 
    bar[width=4](0)
 
    var data = DTypePointer[dtype].alloc(tensor.num_elements())
    memcpy(data, tensor.data(), tensor.num_elements())
    print("fourth", data.load[width=4](0))
 
 
fn main():
    var h = rand[dtype](3, 3)
 
    print("\n owned tensor")
    foo(h)
    # foo(h ^) same behavior
 
    print("\n owned tensor 2")
    foo_2(h)
 
    _ = h

values printed

owned tensor
first 0.1315377950668335
second [0.1315377950668335, 0.458650141954422, 0.21895918250083923, 0.67886471748352051]
third [4.2876297505565036e-32, 3.1799666050923074e-41, 0.21895918250083923, 0.67886471748352051]
 
owned tensor 2
first 0.1315377950668335
second [0.1315377950668335, 0.458650141954422, 0.21895918250083923, 0.67886471748352051]
third [0.1315377950668335, 0.458650141954422, 0.21895918250083923, 0.67886471748352051]
fourth [0.0, 0.0, 0.21895918250083923, 0.67886471748352051]

System information

- What OS did you do install Mojo on ? Pop.os
- Provide version information for Mojo by pasting the output of `mojo -v` mojo 24.3.0 (9882e19d)
- Provide Modular CLI version by pasting the output of `modular -v` modular 0.7.4 (df7a9e8b)

[Benny-Nottonson]

Slightly smaller and easier to read example

from tensor import Tensor, TensorShape
from testing import assert_true

fn move_tensor[dtype: DType](data: DTypePointer[dtype], owned tensor: Tensor[dtype]):
    memcpy(data, tensor.data(), tensor.num_elements())
    # NOTE: adding the line `_ = tensor` solves the issue, but it should not be necessary

fn main():
    alias shape = TensorShape(3, 3, 3)
    alias dtype = DType.float32

    var tensor = Tensor[dtype](shape)
    var data = DTypePointer[dtype].alloc(shape.num_elements())

    for i in range(shape.num_elements()):
        tensor[i] = i
    
    move_tensor(data, tensor)

    for i in range(shape.num_elements()):
        try:
            assert_true(data[i] == i)
        except:
            print("data[", i, "] = ", data[i], " != ", i)

[Benny-Nottonson]

The issue only occurs when 1. The tensor is passed in as owned and 2. The tensor is not explicitly deleted in the move function. Also reproduces on WSL2 Ubuntu and MacOS, so it seems independent of system