[PIR] pir onednn support conv3d,conv3d_grad (#61062)

wanghuancoder · web-flow · commit b771e3f050e7 · 2024-01-25T13:59:46.000+08:00
* pir onednn support conv3d,conv3d_grad
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_instruction.cc
@@ -344,6 +344,11 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(
     }
   }
   TensorNameMap(op, *value_exec_info_, yaml_info_parser, inputs_, outputs_);
+
+  // Step4: Mark is_run_mkldnn_kernel=true
+  phi::MetaConfig new_config = infer_meta_context_.GetMetaConfig();
+  new_config.is_run_mkldnn_kernel = true;
+  infer_meta_context_.SetMetaConfig(new_config);
 }
 
 OneDNNPhiKernelInstruction::~OneDNNPhiKernelInstruction() {
diff --git a/paddle/fluid/framework/new_executor/instruction/onednn/onednn_legacy_instruction.cc b/paddle/fluid/framework/new_executor/instruction/onednn/onednn_legacy_instruction.cc
@@ -226,6 +226,11 @@ OneDNNLegacyKernelInstruction::OneDNNLegacyKernelInstruction(
           op_normalizer.GetLegacyArgName(fluid_op_name, input_name));
     }
   }
+
+  // Step3: Mark is_run_mkldnn_kernel=true
+  phi::MetaConfig new_config = infer_meta_context_.GetMetaConfig();
+  new_config.is_run_mkldnn_kernel = true;
+  infer_meta_context_.SetMetaConfig(new_config);
 }
 
 OneDNNLegacyKernelInstruction::~OneDNNLegacyKernelInstruction() {
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -19,6 +19,14 @@
   extra_args : bool is_test=false
   data_format_tensors : input, out_grad
 
+- op : conv3d
+  extra_args : bool is_test=false
+  data_format_tensors : input
+
+- op : conv3d_grad
+  extra_args : bool is_test=false
+  data_format_tensors : input, out_grad
+
 - op : fused_conv2d
   extra_args : float fuse_alpha = 0.0, float fuse_beta = 0.0, float scale_in=1.0, float scale_out=1.0, float scale_in_eltwise=1.0, float[] scale_weights={1.0f}
   data_format_tensors : input
diff --git a/test/ir/inference/test_mkldnn_conv3d_op.py b/test/ir/inference/test_mkldnn_conv3d_op.py
@@ -17,7 +17,7 @@
 
 import hypothesis.strategies as st
 import numpy as np
-from auto_scan_test import MkldnnAutoScanTest
+from auto_scan_test import MkldnnAutoScanTest, PirMkldnnAutoScanTest
 from hypothesis import given
 from program_config import OpConfig, ProgramConfig, TensorConfig
 
@@ -91,5 +91,75 @@ def test(self, *args, **kwargs):
         self.run_test(*args, **kwargs)
 
 
+class TestPirOneDNNPad3DOp(PirMkldnnAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self, *args, **kwargs):
+        def generate_input(*args, **kwargs):
+            if kwargs["data_format"] == "NCDHW":
+                return np.random.random(
+                    [kwargs["batch_size"], 48, 64, 32, 64]
+                ).astype(np.float32)
+            else:
+                return np.random.random(
+                    [kwargs["batch_size"], 64, 32, 64, 48]
+                ).astype(np.float32)
+
+        def generate_weight(*args, **kwargs):
+            return np.random.random(
+                [16, int(48 / kwargs["groups"]), 3, 3, 3]
+            ).astype(np.float32)
+
+        conv3d_op = OpConfig(
+            type="conv3d",
+            inputs={"Input": ["input_data"], "Filter": ["conv_weight"]},
+            outputs={"Output": ["conv_output"]},
+            attrs={
+                "data_format": kwargs["data_format"],
+                "dilations": kwargs["dilations"],
+                "padding_algorithm": kwargs["padding_algorithm"],
+                "groups": kwargs["groups"],
+                "paddings": kwargs["paddings"],
+                "strides": kwargs["strides"],
+                "is_test": True,
+                "use_mkldnn": True,
+            },
+        )
+
+        program_config = ProgramConfig(
+            ops=[conv3d_op],
+            weights={
+                "conv_weight": TensorConfig(
+                    data_gen=partial(generate_weight, *args, **kwargs)
+                )
+            },
+            inputs={
+                "input_data": TensorConfig(
+                    data_gen=partial(generate_input, *args, **kwargs)
+                )
+            },
+            outputs=["conv_output"],
+        )
+
+        yield program_config
+
+    def sample_predictor_configs(self, program_config):
+        config = self.create_inference_config(use_mkldnn=True)
+        yield config, (1e-5, 1e-5)
+
+    @given(
+        data_format=st.sampled_from(["NCDHW", "NDHWC"]),
+        dilations=st.sampled_from([[1, 2, 1]]),
+        padding_algorithm=st.sampled_from(["EXPLICIT"]),
+        groups=st.sampled_from([2]),
+        paddings=st.sampled_from([[0, 3, 2]]),
+        strides=st.sampled_from([[1, 2, 1]]),
+        batch_size=st.integers(min_value=1, max_value=4),
+    )
+    def test(self, *args, **kwargs):
+        self.run_test(*args, **kwargs)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/legacy_test/test_conv3d_op.py b/test/legacy_test/test_conv3d_op.py
@@ -208,7 +208,10 @@ def init_kernel_type(self):
         def test_check_output(self):
             place = core.CUDAPlace(0)
             self.check_output_with_place(
-                place, check_dygraph=(not self.use_mkldnn), check_pir=True
+                place,
+                check_dygraph=(not self.use_mkldnn),
+                check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
         def test_check_grad_no_filter(self):
@@ -223,6 +226,7 @@ def test_check_grad_no_filter(self):
                 check_dygraph=(not self.use_mkldnn),
                 user_defined_grads=[numeric_grads],
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
         def test_check_grad_no_input(self):
@@ -237,6 +241,7 @@ def test_check_grad_no_input(self):
                 check_dygraph=(not self.use_mkldnn),
                 user_defined_grads=[numeric_grads],
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
         def test_check_grad(self):
@@ -251,6 +256,7 @@ def test_check_grad(self):
                 user_defined_grads=[numeric_input_grads, numeric_fliter_grads],
                 check_dygraph=(not self.use_mkldnn),
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
     cls_name = "{}_{}".format(parent.__name__, "CUDNNBF16OP")
@@ -455,6 +461,7 @@ def test_check_output(self):
             atol=1e-5,
             check_dygraph=(not self.use_mkldnn),
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad(self):
@@ -467,6 +474,7 @@ def test_check_grad(self):
             max_relative_error=0.03,
             check_dygraph=(not self.use_mkldnn),
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_filter(self):
@@ -480,6 +488,7 @@ def test_check_grad_no_filter(self):
             no_grad_set={'Filter'},
             check_dygraph=(not self.use_mkldnn),
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_input(self):
@@ -493,6 +502,7 @@ def test_check_grad_no_input(self):
             no_grad_set={'Input'},
             check_dygraph=(not self.use_mkldnn),
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def init_test_case(self):
@@ -608,7 +618,12 @@ def test_check_output(self):
         if core.is_compiled_with_cuda():
             place = core.CUDAPlace(0)
             if core.is_float16_supported(place):
-                self.check_output_with_place(place, atol=2e-2, check_pir=True)
+                self.check_output_with_place(
+                    place,
+                    atol=2e-2,
+                    check_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
+                )
 
 
 @unittest.skipIf(
@@ -632,7 +647,12 @@ def test_check_output(self):
         if core.is_compiled_with_cuda():
             place = core.CUDAPlace(0)
             if core.is_float16_supported(place):
-                self.check_output_with_place(place, atol=2e-2, check_pir=True)
+                self.check_output_with_place(
+                    place,
+                    atol=2e-2,
+                    check_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
+                )
 
 
 @unittest.skipIf(
@@ -656,7 +676,12 @@ def test_check_output(self):
         if core.is_compiled_with_cuda():
             place = core.CUDAPlace(0)
             if core.is_float16_supported(place):
-                self.check_output_with_place(place, atol=2e-2, check_pir=True)
+                self.check_output_with_place(
+                    place,
+                    atol=2e-2,
+                    check_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
+                )
 
 
 @unittest.skipIf(
@@ -680,7 +705,12 @@ def test_check_output(self):
         if core.is_compiled_with_cuda():
             place = core.CUDAPlace(0)
             if core.is_float16_supported(place):
-                self.check_output_with_place(place, atol=2e-2, check_pir=True)
+                self.check_output_with_place(
+                    place,
+                    atol=2e-2,
+                    check_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
+                )
 
 
 @unittest.skipIf(
@@ -704,7 +734,12 @@ def test_check_output(self):
         if core.is_compiled_with_cuda():
             place = core.CUDAPlace(0)
             if core.is_float16_supported(place):
-                self.check_output_with_place(place, atol=2e-2, check_pir=True)
+                self.check_output_with_place(
+                    place,
+                    atol=2e-2,
+                    check_pir=True,
+                    check_pir_onednn=self.check_pir_onednn,
+                )
 
 
 class TestCUDNNExhaustiveSearch(TestCUDNN):
@@ -780,7 +815,12 @@ def has_cudnn(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0) if self.has_cudnn() else core.CPUPlace()
-        self.check_output_with_place(place, atol=1e-5, check_pir=True)
+        self.check_output_with_place(
+            place,
+            atol=1e-5,
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+        )
 
     def test_check_grad(self):
         if self.dtype == np.float16:
@@ -792,6 +832,7 @@ def test_check_grad(self):
             'Output',
             max_relative_error=0.03,
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_filter(self):
@@ -805,6 +846,7 @@ def test_check_grad_no_filter(self):
             max_relative_error=0.03,
             no_grad_set={'Filter'},
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def test_check_grad_no_input(self):
@@ -818,6 +860,7 @@ def test_check_grad_no_input(self):
             max_relative_error=0.03,
             no_grad_set={'Input'},
             check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
         )
 
     def init_test_case(self):
diff --git a/test/mkldnn/test_conv3d_mkldnn_op.py b/test/mkldnn/test_conv3d_mkldnn_op.py
@@ -30,48 +30,55 @@ def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestMKLDNNCase1(TestCase1):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestMKLDNNGroup1(TestWithGroup1):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestMKLDNNGroup2(TestWithGroup2):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestMKLDNNWith1x1(TestWith1x1):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestMKLDNNWithInput1x1Filter1x1(TestWithInput1x1Filter1x1):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestConv3DOp_AsyPadding_MKLDNN(TestConv3DOp):
     def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
     def init_paddings(self):
         self.pad = [1, 0, 1, 0, 0, 2]
@@ -87,6 +94,7 @@ def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 class TestConv3DOp_Valid_MKLDNN(TestConv3DOp_AsyPadding_MKLDNN):
@@ -98,6 +106,7 @@ def init_kernel_type(self):
         self.use_mkldnn = True
         self.data_format = "NCHW"
         self.dtype = np.float32
+        self.check_pir_onednn = True
 
 
 if __name__ == '__main__':

Original file line number	Diff line number	Diff line change
`@@ -344,6 +344,11 @@ OneDNNPhiKernelInstruction::OneDNNPhiKernelInstruction(`
`344`	`344`	`}`
`345`	`345`	`}`
`346`	`346`	`TensorNameMap(op, *value_exec_info_, yaml_info_parser, inputs_, outputs_);`
	`347`	`+`
	`348`	`+ // Step4: Mark is_run_mkldnn_kernel=true`
	`349`	`+ phi::MetaConfig new_config = infer_meta_context_.GetMetaConfig();`
	`350`	`+ new_config.is_run_mkldnn_kernel = true;`
	`351`	`+ infer_meta_context_.SetMetaConfig(new_config);`
`347`	`352`	`}`
`348`	`353`
`349`	`354`	`OneDNNPhiKernelInstruction::~OneDNNPhiKernelInstruction() {`
Original file line number	Diff line number	Diff line change
`@@ -226,6 +226,11 @@ OneDNNLegacyKernelInstruction::OneDNNLegacyKernelInstruction(`
`226`	`226`	`op_normalizer.GetLegacyArgName(fluid_op_name, input_name));`
`227`	`227`	`}`
`228`	`228`	`}`
	`229`	`+`
	`230`	`+ // Step3: Mark is_run_mkldnn_kernel=true`
	`231`	`+ phi::MetaConfig new_config = infer_meta_context_.GetMetaConfig();`
	`232`	`+ new_config.is_run_mkldnn_kernel = true;`
	`233`	`+ infer_meta_context_.SetMetaConfig(new_config);`
`229`	`234`	`}`
`230`	`235`
`231`	`236`	`OneDNNLegacyKernelInstruction::~OneDNNLegacyKernelInstruction() {`