[0D-Tensor] Support matmul, fix infershape (#55316)

jiahy0825 · web-flow · commit ce8455c029de · 2023-07-13T11:05:00.000+08:00
diff --git a/paddle/cinn/hlir/pe/transform.cc b/paddle/cinn/hlir/pe/transform.cc
@@ -90,7 +90,8 @@ std::vector<std::vector<int>> GetMatmulNewShapes(
                           : std::vector<int>{1, x_shape[0]};
     new_y_shape = trans_y ? std::vector<int>{1, y_shape[0]}
                           : std::vector<int>{y_shape[0], 1};
-    out_shape = {1};
+    // [m] * [m] -> [], which aligns with Paddle's matmul
+    out_shape = {};
   } else if (x_dim == 1) {
     // vector * matrix
     int y_K = trans_y ? y_shape[max_dim - 1] : y_shape[max_dim - 2];
diff --git a/test/cinn/ops/test_zero_dim_tensor.py b/test/cinn/ops/test_zero_dim_tensor.py
@@ -824,5 +824,50 @@ def init_input(self):
         self.target_shape = ()
 
 
+@OpTestTool.skip_if(
+    not is_compiled_with_cuda(), "x86 test will be skipped due to timeout."
+)
+class TestMatmulOp(OpTest):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.init_input()
+
+    def init_input(self):
+        self.inputs = {
+            "x": np.random.randint(-10, 10, [10]).astype(self.dtype),
+            "y": np.random.randint(-10, 10, [10]).astype(self.dtype),
+        }
+        self.target_shape = ()
+
+    def build_paddle_program(self, target):
+        x = paddle.to_tensor(self.inputs["x"], stop_gradient=False)
+        y = paddle.to_tensor(self.inputs["y"], stop_gradient=False)
+        out = paddle.matmul(x, y)
+
+        self.paddle_outputs = [out]
+
+    def build_cinn_program(self, target):
+        builder = NetBuilder("matmul_op")
+        x = builder.create_input(
+            cinn_dtype_convert(self.dtype), self.inputs["x"].shape, "x"
+        )
+        y = builder.create_input(
+            cinn_dtype_convert(self.dtype), self.inputs["y"].shape, "y"
+        )
+        out = builder.matmul(x, y)
+
+        prog = builder.build()
+        res = self.get_cinn_output(
+            prog, target, [x, y], [self.inputs["x"], self.inputs["y"]], [out]
+        )
+
+        self.cinn_outputs = res
+        self.assertEqual(res[0].shape, self.target_shape)
+
+    def test_check_results(self):
+        self.check_outputs_and_grads()
+
+
 if __name__ == "__main__":
     unittest.main()
