is the tir output as below?
with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
print(tvm.lower(s, [A, B, C], simple_mode=False))
is the tir output as below?
with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
print(tvm.lower(s, [A, B, C], simple_mode=False))