Disambiguate python buffers (#2724)

Author: hunhoffe

programming_examples/basic/tiling_exploration/per_tile/per_tile.py

@@ -7,9 +7,8 @@
 # (c) Copyright 2024 Advanced Micro Devices, Inc. or its affiliates
 import argparse
 import numpy as np
-import sys
 
-from aie.iron import LocalBuffer, ObjectFifo, Program, Runtime, Worker
+from aie.iron import Buffer, ObjectFifo, Program, Runtime, Worker
 from aie.iron.placers import SequentialPlacer
 from aie.iron.device import NPU1Col1
 from aie.iron.controlflow import range_
@@ -39,21 +38,18 @@ def generate_module(
 
     # Use an ObjectFifo for dataflow
     of_out = ObjectFifo(flattened_tile)
+    access_counter = Buffer(initial_value=np.array([0], dtype=dtype))
 
     # The task a core will run
-    def access_order(of_out):
-        access_counter = LocalBuffer(initial_value=np.array([0], dtype=dtype))
-
-        for _ in range_(sys.maxsize):
-            elemOut = of_out.acquire(1)
-            for i in range_(tile_size):
-                elemOut[i] = access_counter[0]
-                access_counter[0] += 1
-            of_out.release(1)
-        pass
+    def access_order(of_out, counter_buf):
+        elemOut = of_out.acquire(1)
+        for i in range_(tile_size):
+            elemOut[i] = counter_buf[0]
+            counter_buf[0] += 1
+        of_out.release(1)
 
     # Create a worker (which will be placed on a core) to run the task
-    worker = Worker(access_order, [of_out.prod()], while_true=False)
+    worker = Worker(access_order, [of_out.prod(), access_counter])
 
     # Runtime operations to move data to/from the AIE-array
     rt = Runtime()

programming_examples/basic/vector_reduce_max/multi_column_designs/col_wise_vector_reduce_max.py

@@ -16,7 +16,7 @@
     Program,
     Runtime,
     Worker,
-    LocalBuffer,
+    Buffer,
     str_to_dtype,
 )
 from aie.iron.placers import SequentialPlacer
@@ -72,30 +72,40 @@ def my_reduce_max(dev, in1_size, out_size, num_cores, dtype_str, trace_size):
         else np.array([np.iinfo(dtype).min], dtype=dtype)
     )
 
-    def core_body(*args):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_num_elements,), np.dtype[dtype]],
-            initial_value=min_val,
+    nextC_buffers = []
+    tmp_buffers = []
+    for i in range(num_cores):
+        nextC_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_num_elements,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_num_elements,), np.dtype[dtype]],
-            initial_value=min_val,
+        tmp_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_num_elements,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
+
+    def core_body(*args):
         # Extract fixed arguments from end of args list
         compute_max = args[-1]
         reduce_max_vector = args[-2]
+        tmp_buffer = args[-3]
+        c_buffer = args[-4]
 
         # Extract object fifos from start of args list
         of_in1 = args[0]
         of_out = args[1]
         neighbor_of_in1s = args[
-            2:-2
+            2:-4
         ]  # Variable number of input fifos based on num_cores
 
         for _ in range_(N_div_n):
             elem_in1 = of_in1.acquire(1)
             reduce_max_vector(elem_in1, tmp_buffer, tile_size)
-            compute_max(nextC_buffer, tmp_buffer, nextC_buffer)
+            compute_max(c_buffer, tmp_buffer, c_buffer)
             of_in1.release(1)
 
         elem_out = of_out.acquire(1)
@@ -107,14 +117,14 @@ def core_body(*args):
 
             # Compute max across all inputs
             for elem in elem_in1s[:-1]:
-                compute_max(elem, nextC_buffer, nextC_buffer)
-            compute_max(elem_in1s[-1], nextC_buffer, elem_out)
+                compute_max(elem, c_buffer, c_buffer)
+            compute_max(elem_in1s[-1], c_buffer, elem_out)
 
             # Release all inputs
             for neighbor_of in neighbor_of_in1s:
                 neighbor_of.release(1)
         else:
-            elem_out[0] = nextC_buffer[0]
+            elem_out[0] = c_buffer[0]
         of_out.release(1)
 
     # Define a worker to run the task on a core
@@ -126,7 +136,9 @@ def core_body(*args):
             if num_cores - cores_per_col < i:
                 fifo_args.append(of_outs[i - 1].cons())
 
-        fifo_args.extend([reduce_max_vector, compute_max])
+        fifo_args.extend(
+            [nextC_buffers[i], tmp_buffers[i], reduce_max_vector, compute_max]
+        )
         my_workers.append(
             Worker(
                 core_body,

programming_examples/basic/vector_reduce_max/multi_column_designs/row_wise_vector_reduce_max.py

@@ -16,7 +16,7 @@
     Program,
     Runtime,
     Worker,
-    LocalBuffer,
+    Buffer,
     str_to_dtype,
 )
 from aie.iron.placers import SequentialPlacer
@@ -73,6 +73,21 @@ def my_reduce_max(dev, in1_size, out_size, dtype_str, trace_size):
         if dtype_str == "bf16"
         else np.array([np.iinfo(dtype).min], dtype=dtype)
     )
+    nextC_buffers = []
+    tmp_buffers = []
+    for i in range(n_cores):
+        nextC_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
+        )
+        tmp_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
+        )
 
     taps = [
         TensorAccessPattern(
@@ -85,22 +100,16 @@ def my_reduce_max(dev, in1_size, out_size, dtype_str, trace_size):
     ]
 
     def core_body(*args):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
         # Extract fixed arguments from end of args list
         compute_max = args[-1]
         reduce_max_vector = args[-2]
+        nextC_buffer = args[-3]
+        tmp_buffer = args[-4]
 
         # Extract object fifos from start of args list
         of_in = args[0]
         of_out = args[1]
-        in_fifos = args[2:-2]  # Variable number of input fifos based on n_cores
+        in_fifos = args[2:-4]  # Variable number of input fifos based on n_cores
 
         for _ in range_(num_iter):
             elem_in = of_in.acquire(1)
@@ -150,7 +159,10 @@ def core_body(*args):
                     fifo_args.append(out_fifos[4].cons())
                     fifo_args.extend(out_fifos[j].cons() for j in range(6, n_cores))
 
-        fifo_args.extend([reduce_max_vector, compute_max])
+        fifo_args.extend(
+            [tmp_buffers[i], nextC_buffers[i], reduce_max_vector, compute_max]
+        )
+
         workers.append(
             Worker(
                 core_body,

programming_examples/basic/vector_reduce_max/single_column_designs/vector_reduce_max_chained.py

@@ -15,7 +15,7 @@
     Program,
     Runtime,
     Worker,
-    LocalBuffer,
+    Buffer,
     str_to_dtype,
 )
 from aie.iron.placers import SequentialPlacer
@@ -84,16 +84,26 @@ def my_reduce_max(dev, in1_size, out_size, dtype_str, trace_size):
         else np.array([np.iinfo(dtype).min], dtype=dtype)
     )
 
-    # Define a task to run
-    def start_core_body(of_in, of_out, reduce_max_vector, compute_max):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
+    nextC_buffers = []
+    tmp_buffers = []
+    for i in range(n_cores):
+        nextC_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
+        tmp_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
+
+    # Define a task to run
+    def start_core_body(
+        of_in, of_out, reduce_max_vector, compute_max, nextC_buffer, tmp_buffer
+    ):
         elem_out = of_out.acquire(1)
         for _ in range_(num_iter):
             elem_in = of_in.acquire(1)
@@ -103,16 +113,9 @@ def start_core_body(of_in, of_out, reduce_max_vector, compute_max):
         elem_out[0] = nextC_buffer[0]
         of_out.release(1)
 
-    def core_body(of_in, of_out, in0, reduce_max_vector, compute_max):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
-
+    def core_body(
+        of_in, of_out, in0, reduce_max_vector, compute_max, nextC_buffer, tmp_buffer
+    ):
         for _ in range_(num_iter):
             elem_in = of_in.acquire(1)
             reduce_max_vector(elem_in, tmp_buffer, elems_per_core)
@@ -138,6 +141,8 @@ def core_body(of_in, of_out, in0, reduce_max_vector, compute_max):
                         out_fifos[i + 1].cons(),
                         reduce_max_vector,
                         compute_max,
+                        nextC_buffers[i],
+                        tmp_buffers[i],
                     ],
                     trace=enable_trace,
                 )
@@ -151,6 +156,8 @@ def core_body(of_in, of_out, in0, reduce_max_vector, compute_max):
                         out_fifos[i].prod(),
                         reduce_max_vector,
                         compute_max,
+                        nextC_buffers[i],
+                        tmp_buffers[i],
                     ],
                     trace=enable_trace,
                 )

programming_examples/basic/vector_reduce_max/single_column_designs/vector_reduce_max_memtile.py

@@ -15,7 +15,7 @@
     Program,
     Runtime,
     Worker,
-    LocalBuffer,
+    Buffer,
     str_to_dtype,
 )
 from aie.iron.placers import SequentialPlacer
@@ -93,16 +93,26 @@ def my_reduce_max(dev, in1_size, out_size, dtype_str, trace_size):
         else np.array([np.iinfo(dtype).min], dtype=dtype)
     )
 
-    # Define a task to run
-    def start_core_body(of_in, of_out, reduce_max_vector, compute_max):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
+    nextC_buffers = []
+    tmp_buffers = []
+    for i in range(n_cores):
+        nextC_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
+        tmp_buffers.append(
+            Buffer(
+                type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
+                initial_value=min_val,
+            )
         )
+
+    # Define a task to run
+    def start_core_body(
+        of_in, of_out, reduce_max_vector, compute_max, nextC_buffer, tmp_buffer
+    ):
         elem_out = of_out.acquire(1)
         for _ in range_(num_iter):
             elem_in = of_in.acquire(1)
@@ -120,15 +130,9 @@ def core_body(
         reduce_max_vector,
         reduce_max_scalar,
         compute_max,
+        nextC_buffer,
+        tmp_buffer,
     ):
-        nextC_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
-        tmp_buffer = LocalBuffer(
-            type=np.ndarray[(out_tensor_size,), np.dtype[dtype]],
-            initial_value=min_val,
-        )
         elem_out = elemC_out.acquire(1)
         for _ in range_(num_iter):
             elem_in = of_in.acquire(1)
@@ -156,6 +160,8 @@ def core_body(
                         out_fifos[i].prod(),
                         reduce_max_vector,
                         compute_max,
+                        nextC_buffers[i],
+                        tmp_buffers[i],
                     ],
                     trace=True if i == 1 else None,
                 )
@@ -169,6 +175,8 @@ def core_body(
                 reduce_max_vector,
                 reduce_max_scalar,
                 compute_max,
+                nextC_buffers[i],
+                tmp_buffers[i],
             ]
             workers.append(Worker(core_body, fn_args=fifo_args, trace=None))