Fix the argument list error and enable the compile of program when cinn is enable through environ variables.

Author: Xreki

api/common/paddle_op_benchmark.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import abc
 import six
 import sys
@@ -109,6 +110,22 @@ def generate_gradients(self, targets, inputs):
         gradients = paddle.static.gradients(targets, inputs)
         return gradients
 
+    def compile(self, program):
+        use_cinn = os.environ.get("FLAGS_use_cinn", False)
+        if use_cinn:
+            # To enable CINN, we need to use CompiledProgram to compile the program.
+            # Only forward ops are enabled because loss_name should not be none when
+            # backward ops are contained in the origin program.
+            build_strategy = paddle.static.BuildStrategy()
+            exec_strategy = paddle.static.ExecutionStrategy()
+            exec_strategy.num_threads = 1
+            compiled_program = paddle.static.CompiledProgram(
+                program).with_data_parallel(
+                    build_strategy=build_strategy, exec_strategy=exec_strategy)
+            return compiled_program
+        else:
+            return program
+
     def init_feed_tensor(self, use_gpu, feed_vars, feed_dict, scope):
         place = paddle.CUDAPlace(0) if use_gpu else paddle.CPUPlace()
         for var in feed_vars:
@@ -426,10 +443,12 @@ def _run_static_impl(self,
         executor = paddle.static.Executor(place)
         executor.run(self.startup_program)
 
+        main_program = self._helper.compile(self.main_program)
+
         def _run_main_iter():
             feed_dict = feed if self._need_feed else None
             fetch_vars = self.fetch_list if self._need_fetch else None
-            outputs = executor.run(program=self.main_program,
+            outputs = executor.run(program=main_program,
                                    feed=feed_dict,
                                    fetch_list=fetch_vars,
                                    use_program_cache=True,

api/common/tensorflow_op_benchmark.py

@@ -198,7 +198,7 @@ def _run_null_graph(self, use_gpu, repeat):
             sess.close()
         return walltimes
 
-    def run_impl(self, use_gpu, feed, repeat=1, profiler="none"):
+    def run_impl(self, use_gpu, config, feed, repeat=1, profiler="none"):
         sess = self._init_session(use_gpu)
 
         #tf.debugging.set_log_device_placement(True)
@@ -301,6 +301,7 @@ def run(self, config, args, use_feed_fetch=True, feeder_adapter=None):
         self.allow_growth = False if args.task == "speed" else True
         outputs, stats = self.run_impl(
             use_gpu=args.use_gpu,
+            config=config,
             feed=feed,
             repeat=args.repeat,
             profiler=args.profiler)