#8112: propagate is b batched for matmul/linear combining

tenstorrent · May 31, 2024 · 121eb3f · 121eb3f
1 parent 539ccba
commit 121eb3f
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git a/ttnn/cpp/pybind11/operations/matmul.hpp b/ttnn/cpp/pybind11/operations/matmul.hpp
@@ -28,7 +28,7 @@ void py_module(py::module& module) {
            const std::optional<const DeviceComputeKernelConfig> compute_kernel_config = std::nullopt,
            const std::optional<const ttnn::CoreGrid> core_grid = std::nullopt) -> ttnn::Tensor {
             return ttnn::operations::matmul::matmul(
-                input_tensor_a, input_tensor_b, /*bias=*/std::nullopt, program_config, memory_config, dtype, activation, compute_kernel_config, core_grid);
+                input_tensor_a, input_tensor_b, /*bias=*/std::nullopt, program_config, memory_config, dtype, activation, compute_kernel_config, core_grid, /*propagate_is_b_batched=*/true);
         },
         py::arg("input_tensor_a"),
         py::arg("input_tensor_b"),

diff --git a/ttnn/cpp/ttnn/operations/matmul.cpp b/ttnn/cpp/ttnn/operations/matmul.cpp
@@ -58,7 +58,8 @@ ttnn::Tensor matmul(
     std::optional<const DataType> dtype,
     const std::optional<const std::string>& activation,
     const std::optional<const DeviceComputeKernelConfig> compute_kernel_config,
-    const std::optional<const ttnn::CoreGrid> core_grid) {
+    const std::optional<const ttnn::CoreGrid> core_grid,
+    const bool propagate_is_b_batched) {
     ttnn::validate_input_tensor("ttnn.linear", input_tensor_a, input_tensor_schemas()[0]);
     ttnn::validate_input_tensor("ttnn.linear", input_tensor_b, input_tensor_schemas()[1]);
     ttnn::validate_input_tensor("ttnn.linear", bias, input_tensor_schemas()[2]);
@@ -92,7 +93,7 @@ ttnn::Tensor matmul(
     }
 
     auto output_tensor = tt::operations::primary::matmul(
-        input_tensor_a, input_tensor_b, post_process_bias ? std::nullopt : bias, program_config, memory_config, dtype, compute_kernel_config, false /*untilize_out*/, user_core_coord, get_fused_activation(activation));
+        input_tensor_a, input_tensor_b, post_process_bias ? std::nullopt : bias, program_config, memory_config, dtype, compute_kernel_config, false /*untilize_out*/, user_core_coord, get_fused_activation(activation), propagate_is_b_batched && input_b_is_batched);
 
     if (post_process_bias) {
         output_tensor = tt::operations::primary::bcast(

diff --git a/ttnn/cpp/ttnn/operations/matmul.hpp b/ttnn/cpp/ttnn/operations/matmul.hpp
@@ -41,7 +41,8 @@ ttnn::Tensor matmul(
     std::optional<const DataType> dtype = std::nullopt,
     const std::optional<const std::string>& activation = std::nullopt,
     const std::optional<const DeviceComputeKernelConfig> compute_kernel_config = std::nullopt,
-    const std::optional<const ttnn::CoreGrid> core_grid = std::nullopt);
+    const std::optional<const ttnn::CoreGrid> core_grid = std::nullopt,
+    const bool propagate_is_b_batched = false);
 
 }  // namespace matmul
 }  // namespace operations