fix BenchGC execution issue when driver is mlir (#337)

Author: xurui1995

test/benchgc/src/benchgc/__main__.py

@@ -275,6 +275,9 @@ def get_module_and_args(flags: argparse.Namespace):
 
     entry = benchgc.mlir.util.get_kernel_func_from_module(module, flags.entry)
 
+    with module.context:
+        entry.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+
     for i, arg in enumerate(args):
         # use zero filling if the arg is return value
         set_default_fill(flags, arg, args, i >= len(entry.type.inputs))

test/benchgc/src/benchgc/mlir/module.py

@@ -41,8 +41,6 @@ def init_module(
                     results=[x.get_mlir_type(ctx) for x in outputs],
                 ),
             )
-            f.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
-
             with ir.InsertionPoint(f.add_entry_block()):
                 block_args = f.entry_block.arguments
                 func.ReturnOp(op_func(ctx, *block_args))
@@ -67,8 +65,6 @@ def init_reduce_module(
                     results=[output.get_mlir_type(ctx)],
                 ),
             )
-            f.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
-
             with ir.InsertionPoint(f.add_entry_block()):
                 block_args = f.entry_block.arguments
 

test/benchgc/src/benchgc/pattern/mlp.py

@@ -104,8 +104,6 @@ def init_module(self, ctx: ir.Context) -> ir.Module:
                         inputs=[src] + weights + bias, results=[result]
                     ),
                 )
-                f.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
-
                 with ir.InsertionPoint(f.add_entry_block()):
                     data = f.entry_block.arguments[0]
                     bias_idx = len(weights) + 1