Skip to content

[TIR] Allow multiple element offsets in Buffer - #10816

Closed
Lunderberg wants to merge 6 commits into
apache:mainfrom
Lunderberg:multiple_elem_offsets
Closed

[TIR] Allow multiple element offsets in Buffer#10816
Lunderberg wants to merge 6 commits into
apache:mainfrom
Lunderberg:multiple_elem_offsets

Conversation

@Lunderberg

Copy link
Copy Markdown
Contributor

Follow-up from #9727. When using Stage.tensorize, the element offset of the array view has one offset for each physical dimension of the buffer. This change replaces BufferNode::elem_offset and BufferNode::offset_factor with arrays. The commits in this PR are separated by category, to help in review.

Updates to Buffer class, to handle multiple elem_offsets

  • Includes size checks to ensure that each parameter that depends on the number of physical dimensions (elem_offsets, offset_factors, and axis_separators) are consistent.

  • Support previous Buffer constructor with a single elem_offset and offset_factor, delegating to the array-based constructor.

Update python calls, with backwards compatibility

  • Buffer.elem_offset and Buffer.offset_factor are implemented as properties, to avoid breakage of calling code.

  • decl_buffer handles single or multiple elem_offset and offset_factor parameters. The parameter name is unchanged, to avoid breaking call sites that use named arguments.

Updated TVMScript parser and printer

Other C++ changes

  • Several other locations that referenced either elem_offsets or offset_factors, which required only minor updates to maintain existing behavior.

Follow-up from apache#9727. When using
`Stage.tensorize`, the element offset of the array view has one offset
for each physical dimension of the buffer. This change replaces
`BufferNode::elem_offset` and `BufferNode::offset_factor` with arrays.
- Includes size checks to ensure that each parameter that depends
on the number of physical dimensions (`elem_offsets`,
`offset_factors`, and `axis_separators`) are consistent.
- Support previous Buffer constructor with a single `elem_offset` and
`offset_factor`, delegating to the array-based constructor.
- `Buffer.elem_offset` and `Buffer.offset_factor` are implemented as
properties, to avoid breakage of calling code.
- `decl_buffer` handles single or multiple `elem_offset` and
`offset_factor` parameters. The parameter name is unchanged, to
avoid breaking call sites that use named arguments.
- Several other locations that referenced either `elem_offsets` or
`offset_factors`, which required only minor updates to maintain
existing behavior.
@Lunderberg
Lunderbergforce-pushed the multiple_elem_offsets branch from 08a65dc to 71d10feCompareMarch 29, 2022 15:34
When constructing a buffer, if `elem_offsets` is either `{}` or
`{0}` (the default values for the two constructors), treat as the
start of the array with correct dimensionality.
@tqchen

tqchen commented Mar 29, 2022

Copy link
Copy Markdown
Member

Thanks @Lunderberg.

The multiple elem_offset is something that ideally we would like to have a deeper thoughts/discussions about. It should be a part of BufferAccess(or addition to) semantics clarification that you and @vinx13 are working on.

Moving from a single elem_offset to multiple is likely a big change and would also need some systematic thinking through the buffer access semantics.

It would also be good to think about alternatives. In TensorIR there is a subregion_match primitive to create a subregion view, and depending on the lower-level translation we may not need a buffer with multiple elem_offset explicitly.

@Lunderberg

Copy link
Copy Markdown
ContributorAuthor

Thank you, and I agree that this type of change requires a larger discussion. My intent was for this draft PR to be an option in that discussion, not a ready-to-merge change on its own, and I should have edited the commit messages to make that intention clearer.

@vinx13 and I also had a brief chat on alternatives, and elem_offset may be avoidable altogether. I think the current use cases could be replaced with match_buffer/buffer_bind_scope, depending on whether a schedule is implemented in Schedulable TIR or in TE, along with using buffer.data of the subregion.

@tqchen

Copy link
Copy Markdown
Member

Thank you @Lunderberg ! yes love all the discussions so far

@areuschareusch added needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it and removed needs-triage PRs or issues that need to be investigated by maintainers to find the right assignees to address it labels Oct 19, 2022
@Lunderberg

Copy link
Copy Markdown
ContributorAuthor

Closing this PR, unlikely to take this development path. Current behavior in which elem_offset is valid only for flat memory spaces has been working well.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@Lunderberg@tqchen@areusch