cornell-zhang/hcl-dialect

[Frontend] Add axis to the update stage

Closed this issue · 1 comments

zzzDavid commented

Related test: mlir/test_dsl_def.py::test_module_declarative_compute_at

        @hcl.def_([a.shape, b.shape, c.shape])
        def add(a, b, c):
            d = hcl.compute(a.shape, lambda *x: a[x] + b[x], "d")
            hcl.update(c, lambda *x: d[x] + 1, "u") # this stage `u` has an empty axis list