You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/05/10 21:33:20 UTC
[GitHub] [tvm] Lunderberg opened a new pull request, #11269: [Draft][TIR][Schedule] Transform layout quality of life
Lunderberg opened a new pull request, #11269:
URL: https://github.com/apache/tvm/pull/11269
Draft PR, seeing whether some of the TE-specific interface features can be pulled over into the TIR-based schedules. These are implemented as `Schedule.transform_layout_sugared`, which wraps around the existing `Schedule.transform_layout`.
* Specify buffer by name, rather than by access_type/index.
* Optional specify block by name, rather than needing a separate call to `Schedule.get_block`.
* Enable `*args` in the index mapping function.
Additional things to add:
* Allow `AXIS_SEPARATOR` in the index mapping function, delegate to `Schedule.set_axis_separators`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r874835692
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
def transform_layout(
self,
- block: BlockRV,
- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],
Review Comment:
Thank you for the catch, and I'll update `TransformLayoutTraits::UnpackedAsPython` with the updated arguments. Is there somewhere that should have an additional test for round-tripping the python conversion? It looks like this `UnpackedAsPython` is used in [`Trace.as_python`](https://github.com/apache/tvm/blob/main/python/tvm/tir/schedule/trace.py#L189), but I don't see any tests that call it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1124025550
Something sort of like the changes here are what I'm picturing. Not complete, because `Union[Tuple[int,str],str]` breaks the `@type_checked` annotation, but something sort of like this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r878384725
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
Review Comment:
Thank you, and re-eanbled following #11289.
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2182,34 +2268,60 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) ->
C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
"""
+ block = self._normalize_block_arg(block)
+ buffer_index_type, buffer_index, buffer_obj = self._normalize_buffer_arg(block, buffer)
+
+ ndim = len(buffer_obj.shape)
if callable(index_map):
- index_map = IndexMap.from_func(index_map)
- assert buffer_index_type in ["read", "write"], "Invalid buffer_index_type"
+ index_map, axis_separators = IndexMap.from_func_with_separators(index_map, ndim=ndim)
+ else:
+ axis_separators = []
+
buffer_index_type_enum = 0 if buffer_index_type == "read" else 1
_ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
self, block, buffer_index, buffer_index_type_enum, index_map
)
+ if axis_separators:
+ _ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_index_type_enum, axis_separators
+ )
- @type_checked
+ # @type_checked
Review Comment:
Thank you, and re-eanbled following #11289.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1127700944
CI functionality tests all passed, updates to resolve sphinx documentation tests, and rebase to resolve merge conflict.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1136013252
CI failure in `Docs: GPU` looks like it was caused by a network timeout. Re-triggering CI.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1137429363
CI failure this time was due to CUDA version in `Docs: GPU`. Looks like the docker image was recently updating, merging main into the dev branch and retriggering CI.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1125075556
#11289 should resolve the type-checking annotation issue with Tuples, if we want to go forward with the other changes here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1133276452
Current CI failures due to Scipy URL change, should restart CI after #11399 lands.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1137753230
Retriggering CI following #11456.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r878367921
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
Review Comment:
enable type checking
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2182,34 +2268,60 @@ def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) ->
C[vi, vj] = B[vi // 16, vj // 16, vi % 16, vj % 16] + 1.0
"""
+ block = self._normalize_block_arg(block)
+ buffer_index_type, buffer_index, buffer_obj = self._normalize_buffer_arg(block, buffer)
+
+ ndim = len(buffer_obj.shape)
if callable(index_map):
- index_map = IndexMap.from_func(index_map)
- assert buffer_index_type in ["read", "write"], "Invalid buffer_index_type"
+ index_map, axis_separators = IndexMap.from_func_with_separators(index_map, ndim=ndim)
+ else:
+ axis_separators = []
+
buffer_index_type_enum = 0 if buffer_index_type == "read" else 1
_ffi_api.ScheduleTransformLayout( # type: ignore # pylint: disable=no-member
self, block, buffer_index, buffer_index_type_enum, index_map
)
+ if axis_separators:
+ _ffi_api.ScheduleSetAxisSeparator( # type: ignore # pylint: disable=no-member
+ self, block, buffer_index, buffer_index_type_enum, axis_separators
+ )
- @type_checked
+ # @type_checked
Review Comment:
enable type checking
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1122896480
@vinx13 A draft of some ease-of-use changes I'd been experimenting with this afternoon. These are python-only, and unpack to the same C++ calls, so it (hopefully) shouldn't cause issues elsewhere.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1130216044
All tests passing, including tests for the `UnpackAsPython`. Marking the PR as ready for review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r874899494
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
def transform_layout(
self,
- block: BlockRV,
- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],
Review Comment:
It should be corrected now, including updates to `verify_trace_roundtrip` to validate the python roundtrip by default. If this finds other breakages unrelated to this PR, I'll limit the python roundtrip checks to opt-in tests, with a tracking issue to resolve any other breakages.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
vinx13 commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1122933594
This is super cool, thanks for the PR. Shall we consider the naming with details of what's being sugared. For example, `axis_separators` are related to physical layout, do you think something like `transform_to_physical_layout` would be more meaningful? Feel free to come up with better names
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on a diff in pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r874873432
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
def transform_layout(
self,
- block: BlockRV,
- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],
Review Comment:
Found it. The [`verify_trace_roundtrip`](https://github.com/apache/tvm/blob/main/python/tvm/tir/schedule/testing.py#L25) helper function verifies that the trace can be round-tripped through JSON, and that the trace can be converted to python, but doesn't verify that the trace can be round-tripped through python.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 merged pull request #11269: [TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
vinx13 merged PR #11269:
URL: https://github.com/apache/tvm/pull/11269
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Lunderberg commented on pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
Lunderberg commented on PR #11269:
URL: https://github.com/apache/tvm/pull/11269#issuecomment-1123861668
I'd definitely be up for a better name, as `transform_layout_sugared` definitely leaves the reader uncertain what is being done. My one concern with having a completely different name is that it may not be clear to a reader that this is the same underlying schedule primitive in C++.
Another alternative would be to merge this as the input handling of `Schedule.transform_layout`. There's already some handling, such as the conversion from `Callable` to `IndexMap`, so this would be similar python-only handling. This would also avoid the issue of having multiple functions that aren't immediately obvious as wrappers around each other.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r874317075
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
def transform_layout(
self,
- block: BlockRV,
- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],
Review Comment:
this will break trace to python conversion. It is implemented here https://github.com/apache/tvm/blob/main/src/tir/schedule/primitive/layout_transformation.cc#L288. Unfortunately CI didn’t catch this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #11269: [Draft][TIR][Schedule] Transform layout quality of life
Posted by GitBox <gi...@apache.org>.
vinx13 commented on code in PR #11269:
URL: https://github.com/apache/tvm/pull/11269#discussion_r874317075
##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -2114,25 +2114,111 @@ def after_unannotate(a: T.handle, b: T.handle) -> None:
########## Schedule: Layout transformation ##########
- @type_checked
+ def _normalize_block_arg(self, block: Union[BlockRV, str]) -> BlockRV:
+ if isinstance(block, str):
+ return self.get_block(block)
+
+ return block
+
+ def _normalize_buffer_arg(
+ self, block: BlockRV, buffer: Union[Tuple[str, int], str, Buffer]
+ ) -> Tuple[str, int, Buffer]:
+
+ block_name = self.get(block).name_hint
+
+ def iter_buffers():
+ block_obj = self.get(block)
+ for i, read in enumerate(block_obj.reads):
+ yield "read", i, read.buffer
+ for i, write in enumerate(block_obj.writes):
+ yield "write", i, write.buffer
+
+ if isinstance(buffer, str):
+ possible_buffers = {}
+ # String lookup requires ensuring that the name is unique
+ for buffer_index, buffer_index_type, buf in iter_buffers():
+ if buf.name == buffer:
+ possible_buffers[buf] = (buffer_index_type, buffer_index)
+
+ assert possible_buffers, f"Could not find buffer '{buffer}' in block '{block_name}'"
+ assert (
+ len(possible_buffers) == 1
+ ), f"Multiple buffers named '{buffer}' in block '{block_name}'"
+ buffer_obj, (buffer_index, buffer_index_type) = next(iter(possible_buffers.items()))
+
+ elif isinstance(buffer, Buffer):
+ # Buffer lookup has unique id, can break out early
+ found = False
+ for buffer_index, buffer_index_type, buffer_obj in iter_buffers():
+ if buffer_obj.same_as(buffer):
+ found = True
+ break
+
+ assert found, "Could not find buffer '{buffer.name}' in block '{block_name}'"
+
+ elif isinstance(buffer, tuple):
+ buffer_index_type, buffer_index = buffer
+ assert buffer_index_type in ["read", "write",], (
+ f"Invalid buffer_index_type. "
+ f"Expected 'read' or 'write', "
+ f"but received {buffer_index_type}"
+ )
+ buffer_list = (
+ self.get(block).reads if buffer_index_type == "read" else self.get(block).writes
+ )
+ assert 0 <= buffer_index < len(buffer_list), (
+ f"Invalid buffer_index {buffer_index}. "
+ f"Block {block_name} has only "
+ f"{len(buffer_list)} {buffer_index_type} buffers."
+ )
+ buffer_obj = buffer_list[buffer_index].buffer
+
+ else:
+ raise TypeError(f"Invalid type for argument 'buffer': {type(buffer)}")
+
+ return (buffer_index_type, buffer_index, buffer_obj)
+
+ # @type_checked
def transform_layout(
self,
- block: BlockRV,
- buffer_index: int,
- buffer_index_type: str,
+ block: Union[BlockRV, str],
+ buffer: Union[Tuple[str, int], str, Buffer],
Review Comment:
this will break trace to python conversion. It is implemented here https://github.com/apache/tvm/blob/main/src/tir/schedule/primitive/layout_transformation.cc#L288.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org