diff --git a/conftest.py b/conftest.py index d611f0dd..1b015aa6 100644 --- a/conftest.py +++ b/conftest.py @@ -18,7 +18,8 @@ def aie_context(request): """Create a fresh AIEContext for each test""" verbose_mlir = request.config.option.verbose > 0 - ctx = AIEContext(mlir_verbose=verbose_mlir) + compiler = request.config.getoption("--compiler", default="peano") + ctx = AIEContext(mlir_verbose=verbose_mlir, compiler=compiler) yield ctx aie_utils.DefaultNPURuntime.cleanup() @@ -35,6 +36,12 @@ def pytest_addoption(parser): default=5, help="Number of iterations to run each test for statistics", ) + parser.addoption( + "--compiler", + default="peano", + choices=["peano", "chess"], + help="Kernel compiler: 'peano' (default) or 'chess' (requires Vitis/aietools)", + ) def get_git_commit(): diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index 7b5c00cd..6c4c059d 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -23,7 +23,7 @@ AieccCompilationRule, AieccFullElfCompilationRule, AieccXclbinInstsCompilationRule, - PeanoCompilationRule, + KernelCompilationRule, ArchiveCompilationRule, ) from .fusion import ( diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index d6d17a64..4b8fe92d 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -482,10 +482,13 @@ def generate_mlir(output_artifact, generator): class AieccCompilationRule(CompilationRule): - def __init__(self, build_dir, peano_dir, mlir_aie_dir, *args, **kwargs): + def __init__( + self, build_dir, peano_dir, mlir_aie_dir, use_chess=False, *args, **kwargs + ): self.build_dir = build_dir self.aiecc_path = Path(mlir_aie_dir) / "bin" / "aiecc" self.peano_dir = peano_dir + self.use_chess = use_chess super().__init__(*args, **kwargs) @@ -503,10 +506,20 @@ def compile(self, graph): "-v", "-j1", "--no-compile-host", - "--no-xchesscc", - "--no-xbridge", - "--peano", - str(self.peano_dir), + ] + if self.use_chess: + compile_cmd += [ + "--xchesscc", + "--xbridge", + ] + else: + compile_cmd += [ + "--no-xchesscc", + "--no-xbridge", + "--peano", + str(self.peano_dir), + ] + compile_cmd += [ "--expand-load-pdis", "--generate-full-elf", "--full-elf-name", @@ -547,10 +560,20 @@ def compile(self, graph): "-v", "-j1", "--no-compile-host", - "--no-xchesscc", - "--no-xbridge", - "--peano", - str(self.peano_dir), + ] + if self.use_chess: + compile_cmd += [ + "--xchesscc", + "--xbridge", + ] + else: + compile_cmd += [ + "--no-xchesscc", + "--no-xbridge", + "--peano", + str(self.peano_dir), + ] + compile_cmd += [ "--dynamic-objFifos", ] do_compile_xclbin = mlir_source in mlir_sources_to_xclbins @@ -624,23 +647,24 @@ def _find_tool(name, peano_dir, mlir_aie_dir): ) -class PeanoCompilationRule(CompilationRule): - def __init__(self, peano_dir, mlir_aie_dir, *args, **kwargs): +class KernelCompilationRule(CompilationRule): + """Compile KernelObjectArtifacts using Peano (clang++) or xchesscc.""" + + def __init__(self, peano_dir, mlir_aie_dir, use_chess=False, *args, **kwargs): self.peano_dir = peano_dir self.mlir_aie_dir = mlir_aie_dir + self.use_chess = use_chess super().__init__(*args, **kwargs) def matches(self, artifacts): return any(artifacts.get_worklist(KernelObjectArtifact)) def compile(self, artifacts): - clang_path = Path(self.peano_dir) / "bin" / "clang++" include_path = Path(self.mlir_aie_dir) / "include" worklist = artifacts.get_worklist(KernelObjectArtifact) commands = [] kernel_dir = get_kernel_dir() - target = f"{kernel_dir}-none-unknown-elf" runtime_lib_include_path = ( Path(self.mlir_aie_dir) / "aie_runtime_lib" / kernel_dir.upper() ) @@ -655,29 +679,39 @@ def compile(self, artifacts): raise RuntimeError( "Expected KernelObject dependency to be a C source file" ) - for extra_dep in list(artifact.dependencies)[1:]: - if not isinstance(extra_dep, SourceArtifact): - raise RuntimeError( - "Expected all KernelObject dependencies to be C source files" - ) - - cmd = ( - [ - str(clang_path), - "-O2", - "-std=c++20", - f"--target={target}", - "-Wno-parentheses", - "-Wno-attributes", - "-Wno-macro-redefined", - "-Wno-empty-body", - "-Wno-missing-template-arg-list-after-template-kw", - f"-I{str(include_path)}", - f"-I{str(runtime_lib_include_path)}", - ] - + artifact.extra_flags - + ["-c", source_file.filename, "-o", artifact.filename] - ) + + if self.use_chess: + wrapper_path = Path(self.mlir_aie_dir) / "bin" / "xchesscc_wrapper" + cmd = ( + [ + str(wrapper_path), + kernel_dir, # e.g. "aie2" or "aie2p" + f"-I{str(include_path)}", + f"-I{str(runtime_lib_include_path)}", + ] + + artifact.extra_flags + + ["-c", source_file.filename, "-o", artifact.filename] + ) + else: + clang_path = Path(self.peano_dir) / "bin" / "clang++" + target = f"{kernel_dir}-none-unknown-elf" + cmd = ( + [ + str(clang_path), + "-O2", + "-std=c++20", + f"--target={target}", + "-Wno-parentheses", + "-Wno-attributes", + "-Wno-macro-redefined", + "-Wno-empty-body", + "-Wno-missing-template-arg-list-after-template-kw", + f"-I{str(include_path)}", + f"-I{str(runtime_lib_include_path)}", + ] + + artifact.extra_flags + + ["-c", source_file.filename, "-o", artifact.filename] + ) commands.append(ShellCompilationCommand(cmd)) if artifact.rename_symbols: @@ -693,9 +727,7 @@ def _find_tool(self, name): def _rename_symbols(self, artifact): objcopy_path = self._find_tool("llvm-objcopy") - cmd = [ - objcopy_path, - ] + cmd = [objcopy_path] for old_sym, new_sym in artifact.rename_symbols.items(): cmd += [ "--redefine-sym", diff --git a/iron/common/context.py b/iron/common/context.py index de522cec..0dc1eee4 100644 --- a/iron/common/context.py +++ b/iron/common/context.py @@ -18,6 +18,9 @@ class AIEContext: base_dir: Repository root directory (three levels above this file). build_dir: Directory where compiled artifacts are written. mlir_verbose: Enable verbose MLIR output during compilation. + compiler: Kernel compiler to use: "peano" (default) or "chess". + When "chess", all kernels and aiecc linking use xchesscc. + Requires Vitis/aietools in PATH. """ # Repo root: iron/common/../../.. = three levels up from this file. @@ -25,10 +28,15 @@ class AIEContext: build_dir: Path = field(default_factory=lambda: Path(os.getcwd()) / "build") mlir_verbose: bool = False + compiler: str = "peano" def __post_init__(self) -> None: """Normalize build_dir to a Path object.""" self.build_dir = Path(self.build_dir) + if self.compiler not in ("peano", "chess"): + raise ValueError( + f"compiler must be 'peano' or 'chess', got {self.compiler!r}" + ) @property def compilation_rules(self): @@ -40,13 +48,17 @@ def compilation_rules(self): """ mlir_aie_dir = Path(aie.utils.config.root_path()) peano_dir = Path(aie.utils.config.peano_install_dir()) + use_chess = self.compiler == "chess" + return [ comp.FusePythonGeneratedMLIRCompilationRule(), comp.GenerateMLIRFromPythonCompilationRule(), - comp.PeanoCompilationRule(peano_dir, mlir_aie_dir), + comp.KernelCompilationRule(peano_dir, mlir_aie_dir, use_chess=use_chess), comp.ArchiveCompilationRule(peano_dir, mlir_aie_dir), comp.AieccXclbinInstsCompilationRule( - self.build_dir, peano_dir, mlir_aie_dir + self.build_dir, peano_dir, mlir_aie_dir, use_chess=use_chess + ), + comp.AieccFullElfCompilationRule( + self.build_dir, peano_dir, mlir_aie_dir, use_chess=use_chess ), - comp.AieccFullElfCompilationRule(self.build_dir, peano_dir, mlir_aie_dir), ] diff --git a/iron/operators/channeled_unary_design.py b/iron/operators/channeled_unary_design.py index 693d2baa..5b8b2d72 100644 --- a/iron/operators/channeled_unary_design.py +++ b/iron/operators/channeled_unary_design.py @@ -71,9 +71,6 @@ def core_fn(of_in, of_out, kernel_line): of_out.release(1) # Create a worker to perform the task - # Large tile sizes (>4096) with LUT-based kernels need more stack space - # than the default 1024 bytes due to spilled vector temporaries. - worker_kwargs = {"stack_size": 0xD00} if line_size > 4096 else {} my_workers = [ Worker( core_fn, @@ -82,7 +79,6 @@ def core_fn(of_in, of_out, kernel_line): of_outs[i * num_channels + j].prod(), kernel_fcn, ], - **worker_kwargs, ) for i in range(num_columns) for j in range(num_channels)