[GitHub] [tvm] manupa-arm commented on a change in pull request #9065: Move the allocates of AoT codegen to be TVMBAWs

[GitBox <gi...@apache.org>]

manupa-arm commented on a change in pull request #9065:
URL: https://github.com/apache/tvm/pull/9065#discussion_r714157805



##########
File path: src/tir/transforms/lower_tvm_builtin.cc
##########
@@ -115,7 +115,8 @@ class BuiltinLower : public StmtExprMutator {
     int64_t nbytes = GetVectorBytes(op->dtype);
     if (device_type_.defined()) {

Review comment:
       Done

##########
File path: src/tir/transforms/storage_rewrite.cc
##########
@@ -478,6 +478,10 @@ class StoragePlanRewriter : public StmtExprMutator {
     uint64_t bits_offset{0};
   };
 
+  bool IsSpecialTaggedMemory(const StorageScope& scope) {

Review comment:
       Done

##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -589,5 +590,41 @@ def test_memory_planning(workspace_byte_alignment, main_workspace_size, sum_work
     )
 
 
+def test_aot_codegen_backend_alloc_workspace_calls():
+    dtype = "float32"
+
+    # These shapes should create small tensors that would
+    # get lowered to stack allocations in the CPU PrimFuncs.
+    # However, the AoT executor codegen should retain them
+    # as TVMBAW calls
+    ishape = (1, 4, 4, 4)
+    wshape = (4, 4, 3, 3)
+
+    data0 = relay.var("data", shape=ishape, dtype=dtype)
+    weight0 = relay.var("weight", shape=wshape, dtype=dtype)
+    out = relay.nn.conv2d(data0, weight0, kernel_size=(3, 3), padding=(1, 1), groups=1)
+    main_f = relay.Function([data0, weight0], out)
+    mod = tvm.IRModule()
+    mod["main"] = main_f
+    mod = transform.InferType()(mod)
+
+    i_data = np.random.uniform(0, 1, ishape).astype(dtype)
+    w1_data = np.random.uniform(0, 1, wshape).astype(dtype)
+
+    inputs = OrderedDict([("data", i_data), ("weight", w1_data)])
+    output_list = generate_ref_data(mod, inputs)
+
+    compiled_runtime_modules = compile_models(

Review comment:
       Ah it is a bit cumbersome to do that :), Instead I used relay in primitive form so its clear that main function should only have three allocates.

##########
File path: tests/python/relay/aot/test_crt_aot.py
##########
@@ -589,5 +590,41 @@ def test_memory_planning(workspace_byte_alignment, main_workspace_size, sum_work
     )
 
 
+def test_aot_codegen_backend_alloc_workspace_calls():
+    dtype = "float32"
+
+    # These shapes should create small tensors that would
+    # get lowered to stack allocations in the CPU PrimFuncs.
+    # However, the AoT executor codegen should retain them
+    # as TVMBAW calls
+    ishape = (1, 4, 4, 4)
+    wshape = (4, 4, 3, 3)
+
+    data0 = relay.var("data", shape=ishape, dtype=dtype)
+    weight0 = relay.var("weight", shape=wshape, dtype=dtype)
+    out = relay.nn.conv2d(data0, weight0, kernel_size=(3, 3), padding=(1, 1), groups=1)
+    main_f = relay.Function([data0, weight0], out)
+    mod = tvm.IRModule()
+    mod["main"] = main_f
+    mod = transform.InferType()(mod)
+
+    i_data = np.random.uniform(0, 1, ishape).astype(dtype)
+    w1_data = np.random.uniform(0, 1, wshape).astype(dtype)
+
+    inputs = OrderedDict([("data", i_data), ("weight", w1_data)])
+    output_list = generate_ref_data(mod, inputs)
+
+    compiled_runtime_modules = compile_models(
+        AOTTestModel(module=mod, inputs=inputs, outputs=output_list),
+        "c",
+        True,
+    )
+
+    source = compiled_runtime_modules[0].lib.imported_modules[0].get_source()

Review comment:
       I can understand the reasoning but the current flow just creates per target IRModules just before runtime.Modules are created. Therefore all host_target (i.e. CPU) PrimFuncs end up in a single runtime.Module. 




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