cornell-zhang / heterocl

HeteroCL: A Multi-Paradigm Programming Infrastructure for Software-Defined Heterogeneous Computing
https://cornell-zhang.github.io/heterocl/
Apache License 2.0
326 stars 92 forks source link

Adding support for getting/setting tensor slices #382

Open seanlatias opened 3 years ago

seanlatias commented 3 years ago

We should add support for the following features.

  1. Get a slice of a tensor and make a copy
    A = hcl.placeholder((10,))
    B = A[2:5].copy()
    # equivalent API
    B = hcl.compute((3,), lambda x: A[x+2])
  2. Set a slice of a tensor
    A = hcl.placeholder((10,))
    A[2:5] = [1, 2, 3]
    A[2:5] = 4  # broadcast
  3. Get a slice of a tensor without making a copy
    A = hcl.placeholder((10, 10))
    B = A[2:5]
    B[0] = 1
    hcl.print(A[2]) # returns 1