Shape of TensorSlice is not correct when indexed by a struct field

[jcasas00]

The following code passes using the main branch but not hcl-mlir.

def test_tensor_slice_struct():
    hcl.init()
    def kernel():
        stype = hcl.Struct({"x": hcl.UInt(8), "y": hcl.UInt(8)})    # x is in the lower bits ...
        xy = hcl.scalar(0x0102, "foo", dtype=stype).v
        z1 = hcl.compute((2,3), lambda x,y: x+y, dtype=hcl.UInt(32))
        r = hcl.compute((4,), lambda _: 0, dtype=hcl.UInt(32))
        t = z1[xy.y]
        assert t.shape == (3,)      # this assertion fails when using hcl-mlir branch
        r[0] = t[xy.x]
        return r
    
    s = hcl.create_schedule([], kernel)
    hcl_res = hcl.asarray(np.zeros((4,), dtype=np.uint32), dtype=hcl.UInt(32))
    f = hcl.build(s)
    f(hcl_res)
    np_res = hcl_res.asnumpy()
    golden = np.zeros((4,), dtype=np.uint32)
    print(f"{np_res}")
    golden[0] = 3
    assert np.array_equal(golden, np_res)

[zzzDavid]

TensorSliceOp did not support ExprOp type slice argument before, this is fixed by b7000c0.

Right now this particular test passes, I will close this thread after I add more tests.

[jcasas00]

Thanks. Works for me now as well.

[zzzDavid]

Two tests for tensor slice/indexing with expression are added here:

mlir/test_dsl_basic.py::test_tensor_slice_struct
mlir/test_dsl_basic.py:: test_tensor_index_expr

Commit point: cornell-zhang/heterocl@60eda2b