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/08/14 06:10:33 UTC

[GitHub] [tvm] junrushao1994 commented on a diff in pull request #12422: [TIR] Allow `tir.Buffer` converted to `BufferLoad/BufferRegion` with `__getitem__`

junrushao1994 commented on code in PR #12422:
URL: https://github.com/apache/tvm/pull/12422#discussion_r945234575


##########
python/tvm/tir/buffer.py:
##########
@@ -176,6 +177,40 @@ def offset_of(self, indices):
         """
         return _ffi_api.BufferOffsetOf(self, indices)  # type: ignore
 
+    def __getitem__(self, indices):
+        from ..arith import Analyzer  # pylint: disable=import-outside-toplevel
+        from .expr import BufferLoad, Ramp  # pylint: disable=import-outside-toplevel
+        from .stmt import BufferRegion  # pylint: disable=import-outside-toplevel
+
+        if not isinstance(indices, (tuple, list)):
+            indices = [indices]
+        if any(isinstance(index, slice) and index.step is None for index in indices):
+            region = []
+            for index in indices:
+                if isinstance(index, slice):
+                    region.append(
+                        Range.from_min_extent(
+                            index.start, Analyzer().simplify(index.stop - index.start)

Review Comment:
   nit: do not create `Analyzer` in each loop instance. Please lift its creation outside of the loop



-- 
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