|
9 | 9 |
|
10 | 10 | from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker |
11 | 11 | from aie.iron.placers import SequentialPlacer |
12 | | -from aie.iron.device import NPU1Col1 |
| 12 | +from aie.iron.device import NPU1Col1, NPU2 |
13 | 13 |
|
14 | | -width = 64 |
15 | | -height = 36 |
16 | | -if len(sys.argv) == 3: |
17 | | - width = int(sys.argv[1]) |
18 | | - height = int(sys.argv[2]) |
| 14 | +from aie.extras.dialects.ext import arith |
| 15 | +from aie.helpers.util import np_ndarray_type_get_shape |
| 16 | +from aie.dialects.aie import T |
19 | 17 |
|
20 | | -lineWidth = width |
21 | | -lineWidthInBytes = width * 4 |
22 | | -tensorSize = width * height * 4 # 4 channels |
23 | 18 |
|
24 | | -traceSize = 1024 |
| 19 | +def color_detect(dev, width, height): |
| 20 | + lineWidth = width |
| 21 | + lineWidthInBytes = width * 4 |
| 22 | + tensorSize = width * height * 4 # 4 channels |
25 | 23 |
|
26 | | - |
27 | | -def color_detect(): |
| 24 | + traceSize = 1024 |
28 | 25 |
|
29 | 26 | # Define types |
30 | 27 | line_bytes_ty = np.ndarray[(lineWidthInBytes,), np.dtype[np.uint8]] |
@@ -212,8 +209,20 @@ def or_gray2rgba_and_fn( |
212 | 209 | rt.drain(outOF_L2L3.cons(), O, wait=True) |
213 | 210 |
|
214 | 211 | # Place components (assign them resources on the device) and generate an MLIR module |
215 | | - return Program(NPU1Col1(), rt).resolve_program(SequentialPlacer()) |
216 | | - |
217 | | - |
218 | | -module = color_detect() |
| 212 | + return Program(dev, rt).resolve_program(SequentialPlacer()) |
| 213 | + |
| 214 | + |
| 215 | +try: |
| 216 | + device_name = str(sys.argv[1]) |
| 217 | + if device_name == "npu": |
| 218 | + dev = NPU1Col1() |
| 219 | + elif device_name == "npu2": |
| 220 | + dev = NPU2() |
| 221 | + else: |
| 222 | + raise ValueError("[ERROR] Device name {} is unknown".format(sys.argv[1])) |
| 223 | + width = 64 if (len(sys.argv) != 4) else int(sys.argv[2]) |
| 224 | + height = 36 if (len(sys.argv) != 4) else int(sys.argv[3]) |
| 225 | +except ValueError: |
| 226 | + print("Argument has inappropriate value") |
| 227 | +module = color_detect(dev, width, height) |
219 | 228 | print(module) |
0 commit comments