Skip to content

Commit

Permalink
Support buffer 1d indexing (apache#61)
Browse files Browse the repository at this point in the history
  • Loading branch information
junrushao committed Jul 13, 2022
1 parent f394937 commit 930bf5f
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions python/tvm/script/builder/tir/var.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,11 @@
# specific language governing permissions and limitations
# under the License.
"""TVM Script TIR Buffer"""
from tvm.ir import PrimExpr, Array, Range
from tvm.tir import Var, IntImm, BufferLoad, BufferRegion
from tvm._ffi import register_object as _register_object
from tvm.runtime import Object, DataType
from tvm.ir import Array, PrimExpr, Range
from tvm.runtime import DataType, Object
from tvm.tir import BufferLoad, BufferRegion, IntImm, Var

from . import _ffi_api


Expand Down Expand Up @@ -94,6 +95,8 @@ def buffer_type(self) -> int:
return self.buffer.buffer_type

def __getitem__(self, indices):
if not isinstance(indices, (tuple, list)):
indices = [indices]
if any(isinstance(index, slice) for index in indices):
region = []
for index in indices:
Expand Down

0 comments on commit 930bf5f

Please sign in to comment.