#3330: merge 2 CreateSemaphore functions into 1

Author: TT-billteng

docs/source/tt_metal/apis/host_apis/buffers/CreateSemaphore.rst

@@ -1,6 +1,4 @@
 CreateSemaphore
 ================
 
-.. doxygenfunction:: CreateSemaphore(Program &program, const CoreRange &core_range, uint32_t initial_value)
-
-.. doxygenfunction:: CreateSemaphore(Program &program, const CoreRangeSet &core_range_set, uint32_t initial_value)
+.. doxygenfunction:: CreateSemaphore(Program &program, const std::variant<CoreRange,CoreRangeSet> &core_spec, uint32_t initial_value)

tt_metal/host_api.hpp

@@ -129,26 +129,13 @@ CircularBufferConfig &GetCircularBufferConfig(Program &program, CircularBufferID
  *
  * Return value: Semaphore address (uint32_t)
  *
- * | Argument      | Description                                          | Type                                                  | Valid Range                                              | Required |
- * |---------------|------------------------------------------------------|-------------------------------------------------------|----------------------------------------------------------|----------|
- * | program       | The program to which semaphore will be added to      | Program &                                             |                                                          | Yes      |
- * | core_range    | Range of the Tensix co-ordinates using the semaphore | const CoreRange & (std::pair<CoreCoord, CoreCoord>)   | Pair of logical coords where first coord <= second coord | Yes      |
- * | initial_value | Initial value of the semaphore                       | uint32_t                                              |                                                          | Yes      |
+ * | Argument      | Description                                          | Type                                                      | Valid Range  | Required |
+ * |---------------|------------------------------------------------------|-----------------------------------------------------------|--------------|----------|
+ * | program       | The program to which semaphore will be added to      | Program &                                                 |              | Yes      |
+ * | core_spec     | Range of the Tensix co-ordinates using the semaphore | const std::variant<CoreRange,CoreRangeSet> &              |              | Yes      |
+ * | initial_value | Initial value of the semaphore                       | uint32_t                                                  |              | Yes      |
  */
-uint32_t CreateSemaphore(Program &program, const CoreRange &core_range, uint32_t initial_value);
-
-/**
- * Initializes semaphore on all cores within core range (inclusive). Each core can have up to four 32B semaphores.
- *
- * Return value: Semaphore address (uint32_t)
- *
- * | Argument       | Description                                                 | Type                   | Valid Range                                               | Required |
- * |----------------|-------------------------------------------------------------|------------------------|-----------------------------------------------------------|----------|
- * | program        | The program to which semaphore will be added to             | Program &              |                                                           | Yes      |
- * | core_range_set    | Set of Range of the Tensix co-ordinates using the semaphore | const CoreRangeSet &   | Pairs of logical coords where first coord <= second coord | Yes      |
- * | initial_value  | Initial value of the semaphore                              | uint32_t               |                                                           | Yes      |
- */
-uint32_t CreateSemaphore(Program &program, const CoreRangeSet &core_range_set, uint32_t initial_value);
+uint32_t CreateSemaphore(Program &program, const std::variant<CoreRange,CoreRangeSet> &core_spec, uint32_t initial_value);
 
 /**
 *  Allocates a DRAM or L1 buffer on device

tt_metal/impl/dispatch/command_queue.cpp

@@ -606,8 +606,8 @@ void send_dispatch_kernel_to_device(Device* device) {
             .compile_args = dispatch_compile_args,
             .defines = consumer_defines});
 
-    tt::tt_metal::CreateSemaphore(dispatch_program, {producer_logical_core, producer_logical_core}, 2);
-    tt::tt_metal::CreateSemaphore(dispatch_program, {consumer_logical_core, consumer_logical_core}, 0);
+    tt::tt_metal::CreateSemaphore(dispatch_program, producer_logical_core, 2);
+    tt::tt_metal::CreateSemaphore(dispatch_program, consumer_logical_core, 0);
 
     detail::CompileProgram(device, dispatch_program);
     tt::tt_metal::detail::ConfigureDeviceWithProgram(device, dispatch_program);

tt_metal/impl/program.hpp

@@ -151,8 +151,7 @@ class Program {
     friend void detail::AddKernel(Program &program, Kernel *kernel);
     friend Kernel *detail::GetKernel(const Program &program, KernelID kernel_id);
 
-    friend uint32_t CreateSemaphore(Program &program, const CoreRangeSet &core_range_set, uint32_t initial_value);
-
+    friend uint32_t CreateSemaphore(Program &program, const std::variant<CoreRange,CoreRangeSet> &core_spec, uint32_t initial_value);
     void add_kernel(Kernel *kernel);
     Kernel *get_kernel(KernelID kernel_id) const;
 

tt_metal/tt_metal.cpp

@@ -219,29 +219,37 @@ CircularBufferConfig &GetCircularBufferConfig(Program &program, CircularBufferID
     return detail::GetCircularBuffer(program, circular_buffer_id)->config();
 }
 
-uint32_t CreateSemaphore(Program &program, const CoreRange &core_range, uint32_t initial_value) {
-    return CreateSemaphore(program, CoreRangeSet({core_range}), initial_value);
-}
-
-uint32_t CreateSemaphore(Program &program, const CoreRangeSet &core_range_set, uint32_t initial_value) {
-    std::optional<uint32_t> address;
-    TT_ASSERT(core_range_set.ranges().size() > 0, "Expecting a non-empty CoreRangeSet!");
-    for (auto core_range : core_range_set.ranges()) {
-        auto start_core = core_range.start;
-        auto end_core = core_range.end;
-        TT_ASSERT(start_core == end_core or start_core < end_core && "Invalid core range!");
-        auto addr = get_semaphore_address(program, core_range);
-        if (!address.has_value()) {
-            address = addr;
-        } else {
-            TT_ASSERT(addr == address);
-        }
-    }
-    TT_ASSERT(address.has_value(), "Expecting a valid Semaphore address!");
+uint32_t CreateSemaphore(Program &program, const std::variant<CoreRange,CoreRangeSet> &core_spec, uint32_t initial_value) {
+    return std::visit(
+        [&](auto&& c) -> uint32_t
+        {
+            using T = std::decay_t<decltype(c)>;
+            CoreRangeSet crs({});
+            if constexpr (std::is_same_v<T, CoreRange>) {
+                crs = CoreRangeSet({c});
+            } else{
+                crs = c;
+            }
+            std::optional<uint32_t> address;
+            TT_ASSERT(crs.ranges().size() > 0, "Expecting a non-empty CoreRangeSet!");
+            for (const auto& core_range : crs.ranges()) {
+                CoreCoord start_core = core_range.start;
+                CoreCoord end_core = core_range.end;
+                TT_ASSERT(start_core == end_core or start_core < end_core && "Invalid core range!");
+                auto addr = get_semaphore_address(program, core_range);
+                if (!address.has_value()) {
+                    address = addr;
+                } else {
+                    TT_ASSERT(addr == address);
+                }
+            }
+            TT_ASSERT(address.has_value(), "Expecting a valid Semaphore address!");
 
-    program.add_semaphore(core_range_set, address.value(), initial_value);
+            program.add_semaphore(crs, address.value(), initial_value);
 
-    return address.value();
+            return address.value();
+        },
+        core_spec);
 }
 
 void WriteToDevice(const Buffer &buffer, const std::vector<uint32_t> &host_buffer) {