xqdan commented on a change in pull request #8044:
URL: https://github.com/apache/tvm/pull/8044#discussion_r632896193
##########
File path: python/tvm/driver/build_module.py
##########
@@ -160,16 +173,38 @@ def lower(sch, args, name="main", binds=None,
simple_mode=False):
lower_phase3 = [x[1] for x in add_lower_pass if x[0] > 2]
# Phase 0
- if isinstance(sch, schedule.Schedule):
- mod = form_irmodule(sch, args, name, binds)
+ pass_list = lower_phase0
+ is_legacy_te_schedule: bool = False
+
+ if isinstance(inputs, schedule.Schedule):
+ if args is None:
+ raise ValueError("args must be given for lowering from TE
schedule")
+ mod = form_irmodule(inputs, args, name, binds)
+ is_legacy_te_schedule = True
+ elif isinstance(inputs, PrimFunc):
+ func = inputs.with_attr("global_symbol", name)
+ if pass_ctx.config.get("tir.noalias", True):
+ func = func.with_attr("tir.noalias", True)
+ mod = tvm.IRModule({name: func})
+ elif isinstance(inputs, IRModule):
+ mod = inputs
else:
- mod = sch
+ raise TypeError(
+ f"tvm.lower expected te.Schedule, PrimFunc or IRModule, but got
{type(inputs)}"
+ )
- pass_list = lower_phase0
# Phase 1
+ if is_legacy_te_schedule:
+ pass_list += [
+ tvm.tir.transform.InjectPrefetch(),
+ tvm.tir.transform.StorageFlatten(64, instrument_bound_checkers),
+ ]
pass_list += [
- tvm.tir.transform.InjectPrefetch(),
- tvm.tir.transform.StorageFlatten(64, instrument_bound_checkers),
+ tvm.tir.transform.LowerInitBlock(),
+ tvm.tir.transform.PlanAndUpdateBufferAllocationLocation(),
+ tvm.tir.transform.ConvertBlocksToOpaque(),
+ tvm.tir.transform.CompactBufferAllocation(),
+ tvm.tir.transform.FlattenBuffer(),
Review comment:
For is_legacy_te_schedule flow，do we need to have both StorageFlatten
and FlattenBuffer？
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
