Add sweeps for complex bw_ops: polar, recip, add, mul (#11364)

* #10147: Added sweeps for complex_polar_bw to ttnn * #10147: Added complex_recip_bw sweeps to ttnn * #10147: Added complex_mul_bw sweeps to ttnn * #10147: Add complexadd_bw sweeps to ttnn * #10147: Added support for dims 2, 3, 4 in complex_backward_mul and complex_backward_add YAML files * #10147: Reformating * #10147: Add num-dims: [2, 3, 4] to complex ops: abs, recip, polar * #10147: Reformating
tenstorrent · Aug 14, 2024 · 18dd2d1 · 18dd2d1
1 parent 99428fa
commit 18dd2d1
Show file tree

Hide file tree

Showing 19 changed files with 435 additions and 0 deletions.
diff --git a/tests/tt_eager/python_api_testing/sweep_tests/pytorch_ops.py b/tests/tt_eager/python_api_testing/sweep_tests/pytorch_ops.py
@@ -2067,6 +2067,67 @@ def cos_bw(x, y, *args, **kwargs):
     return in_data.grad
 
 
+def complex_polar_bw(x, y, *args, **kwargs):
+    grad_data = x
+    in_data = y
+    in_data.requires_grad = True
+
+    in_data.retain_grad()
+    pyt_y = torch.polar(in_data.real, in_data.imag)
+    pyt_y.backward(gradient=grad_data)
+
+    grad_real = torch.real(in_data.grad)
+    grad_imag = torch.imag(in_data.grad)
+
+    return torch.complex(grad_real, grad_imag)
+
+
+def complex_recip_bw(x, y, *args, **kwargs):
+    grad_data = x
+    in_data = y
+    in_data.requires_grad = True
+
+    in_data.retain_grad()
+    pyt_y = torch.reciprocal(in_data)
+    pyt_y.backward(gradient=grad_data)
+
+    return in_data.grad
+
+
+def complex_mul_bw(x, y, z, *args, **kwargs):
+    grad_data = x
+    in_data = y
+    other_data = z
+
+    in_data.requires_grad = True
+    other_data.requires_grad = True
+
+    in_data.retain_grad()
+    other_data.retain_grad()
+
+    pyt_y = in_data * other_data
+    pyt_y.backward(gradient=grad_data)
+
+    return [in_data.grad, other_data.grad]
+
+
+def complex_add_bw(x, y, z, *args, **kwargs):
+    grad_data = x
+    in_data = y
+    other_data = z
+
+    in_data.requires_grad = True
+    other_data.requires_grad = True
+
+    in_data.retain_grad()
+    other_data.retain_grad()
+
+    pyt_y = in_data + other_data
+    pyt_y.backward(gradient=grad_data)
+
+    return [in_data.grad, other_data.grad]
+
+
 def global_avg_pool2d(x, *args, **kwargs):
     output_size = (1, 1)
     x = x.to(torch.float32)

diff --git a/tests/ttnn/python_api_testing/sweep_tests/op_map.py b/tests/ttnn/python_api_testing/sweep_tests/op_map.py
@@ -925,6 +925,10 @@
         "tt_op": ttnn_ops.eltwise_unary_div_no_nan,
         "pytorch_op": pytorch_ops.unary_div_no_nan,
     },
+    "complex-polar-bw": {
+        "tt_op": ttnn_ops.complex_polar_bw,
+        "pytorch_op": pytorch_ops.complex_polar_bw,
+    },
     "complex-conj": {
         "tt_op": ttnn_ops.complex_conj,
         "pytorch_op": pytorch_ops.complex_conj,
@@ -941,4 +945,20 @@
         "tt_op": ttnn_ops.complex_recip,
         "pytorch_op": pytorch_ops.complex_recip,
     },
+    "complex-polar-bw": {
+        "tt_op": ttnn_ops.complex_polar_bw,
+        "pytorch_op": pytorch_ops.complex_polar_bw,
+    },
+    "complex-recip-bw": {
+        "tt_op": ttnn_ops.complex_recip_bw,
+        "pytorch_op": pytorch_ops.complex_recip_bw,
+    },
+    "complex-mul-bw": {
+        "tt_op": ttnn_ops.complex_mul_bw,
+        "pytorch_op": pytorch_ops.complex_mul_bw,
+    },
+    "complex-add-bw": {
+        "tt_op": ttnn_ops.complex_add_bw,
+        "pytorch_op": pytorch_ops.complex_add_bw,
+    },
 }
diff --git a/..._tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_backward_complex_polar_test.yaml b/..._tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_backward_complex_polar_test.yaml
@@ -0,0 +1,26 @@
+---
+test-list:
+  - complex-polar-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-shapes: 2
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_polar_bw_sweep.csv
diff --git a/...ing/sweep_tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_complex_polar_test.yaml b/...ing/sweep_tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_complex_polar_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ing/sweep_tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_complex_recip_test.yaml b/...ing/sweep_tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_complex_recip_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ting/sweep_tests/test_configs/ci_sweep_tests_broken/wormhole/ttnn_complex_recip_test.yaml b/...ting/sweep_tests/test_configs/ci_sweep_tests_broken/wormhole/ttnn_complex_recip_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...p_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_add_test.yaml b/...p_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_add_test.yaml
@@ -0,0 +1,27 @@
+---
+test-list:
+  - complex-add-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
+        num-shapes: 3
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc_list
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16", "BFLOAT8_B"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_add_sweep.csv
diff --git a/...p_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_mul_test.yaml b/...p_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_mul_test.yaml
@@ -0,0 +1,27 @@
+---
+test-list:
+  - complex-mul-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
+        num-shapes: 3
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc_list
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16", "BFLOAT8_B"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_mul_sweep.csv
diff --git a/...tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_recip_test.yaml b/...tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_backward_complex_recip_test.yaml
@@ -0,0 +1,26 @@
+---
+test-list:
+  - complex-recip-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-shapes: 2
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_recip_sweep.csv
diff --git a/...ting/sweep_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_complex_abs_test.yaml b/...ting/sweep_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_complex_abs_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ng/sweep_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_complex_recip_test.yaml b/...ng/sweep_tests/test_configs/ci_sweep_tests_working/grayskull/ttnn_complex_recip_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_add_test.yaml b/...ep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_add_test.yaml
@@ -0,0 +1,27 @@
+---
+test-list:
+  - complex-add-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
+        num-shapes: 3
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc_list
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16", "BFLOAT8_B"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_add_sweep.csv
diff --git a/...ep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_mul_test.yaml b/...ep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_mul_test.yaml
@@ -0,0 +1,27 @@
+---
+test-list:
+  - complex-mul-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
+        num-shapes: 3
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc_list
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16", "BFLOAT8_B"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_mul_sweep.csv
diff --git a/..._tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_polar_test.yaml b/..._tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_polar_test.yaml
@@ -0,0 +1,26 @@
+---
+test-list:
+  - complex-polar-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-shapes: 2
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_polar_bw_sweep.csv
diff --git a/..._tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_recip_test.yaml b/..._tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_backward_complex_recip_test.yaml
@@ -0,0 +1,26 @@
+---
+test-list:
+  - complex-recip-bw:
+      shape:
+        start-shape: [1, 1, 32, 64]
+        end-shape: [6, 12, 256, 256]
+        interval: [1, 1, 32, 64]
+        num-shapes: 2
+        num-samples: 128
+        args-sampling-strategy: "all"
+      env:
+        # TT_PCI_DMA_BUF_SIZE: "1048576"
+      datagen:
+        function: gen_rand_complex
+        args:
+          low: -100
+          high: 100
+      comparison:
+        function: comp_pcc
+      args-gen: gen_dtype_layout_device
+      args:
+        data-layout: ["TILE"]
+        data-type: ["BFLOAT16"]
+        buffer-type: ["DRAM", "L1"]
+        out-buffer-type: ["DRAM", "L1"]
+      output-file: complex_bw_recip_sweep.csv
diff --git a/...sting/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_abs_test.yaml b/...sting/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_abs_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ing/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_polar_test.yaml b/...ing/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_polar_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"

diff --git a/...ing/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_recip_test.yaml b/...ing/sweep_tests/test_configs/ci_sweep_tests_working/wormhole/ttnn_complex_recip_test.yaml
@@ -5,6 +5,7 @@ test-list:
         start-shape: [1, 1, 32, 64]
         end-shape: [6, 12, 256, 256]
         interval: [1, 1, 32, 64]
+        num-dims: [2, 3, 4]
         num-shapes: 1
         num-samples: 128
         args-sampling-strategy: "all"