NicolaLancellotti commented on a change in pull request #10344:
URL: https://github.com/apache/tvm/pull/10344#discussion_r835125952
##########
File path: python/tvm/relay/backend/contrib/ethosu/tir/dma.py
##########
@@ -197,18 +197,72 @@ def get_read_params(stmt):
base_address = [get_base_address(index) for index in inner.value.indices]
data_type = inner.buffer.data.type_annotation.element_type.dtype
+
+ def check_rolling_buffer():
+ rolling_buffer = True
+ floor_mod = None
+
+ def _get_rolling_var(stmt):
+ nonlocal rolling_buffer, floor_mod
+
+ if isinstance(stmt, tvm.tir.FloorMod):
+ if floor_mod is not None:
+ rolling_buffer = False
+ elif (
+ isinstance(stmt.b, tvm.tir.expr.IntImm)
+ and isinstance(stmt.a, tvm.tir.expr.Add)
+ and isinstance(stmt.a.a, tvm.tir.expr.Var)
+ and isinstance(stmt.a.b, tvm.tir.expr.IntImm)
+ ):
+ floor_mod = stmt
+ else:
+ rolling_buffer = False
+ elif isinstance(stmt, tvm.tir.FloorDiv):
+ rolling_buffer = False
+
+ tvm.tir.stmt_functor.post_order_visit(inner.value, _get_rolling_var)
+
+ if rolling_buffer and floor_mod is not None:
+ rolling_var = floor_mod.a.a
+ tile_length = floor_mod.b - floor_mod.a.b
+ return rolling_var, tile_length
Review comment:
Done
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
