feat(nkigen): import NKIgen as a workspace member (revision of #55)#58
Open
shaojiex-aws wants to merge 2 commits into
Open
feat(nkigen): import NKIgen as a workspace member (revision of #55)#58shaojiex-aws wants to merge 2 commits into
shaojiex-aws wants to merge 2 commits into
Conversation
…uron#55) NKIgen is the NumPy-trace -> NISA-dialect lowering frontend (formerly "NKIPyKernelGen"). Imported as a top-level workspace member of the nkipy monorepo, replacing the kernelgen/ subfolder previously proposed in aws-neuron#55. Addresses review feedback on aws-neuron#55: - Drop compiler_explorer/, .claude/skills/, internal scratch dirs (issues/, docs/, tools/kernel_agent/, docker/, qwen3_embedding example, Brazil Config). Compiler Explorer wrapper will land separately as a top-level component shared with other backends. - Drop the duplicated nkipy tests under tests/e2e/nkipy_tests/; the top-level /tests folder is the single source of truth. - Drop BIR-sim Mode and tests; device is the golden, matching how nkipy's main backend tests are structured. - Rename Python package nkipy_kernelgen -> nkigen, harness decorator nkipy_kernelgen_test -> nkigen_test. - Make HardwareConstants.h target-aware: replace global MAX_PARTITION_DIM / MAX_FREE_DIM_MATMUL with getSbufNumPartitions(target) / getMatmulFreeDimTileCap(target). The per-target values mirror nki.backends.mlir_tracer.target_info in the public nki Python wheel. infer-layout, canonicalize-partition-dim, and legalize-layout each take a --target= option (default trn2). - Move project metadata into pyproject.toml; setup.py is now just the CMake build extension. Wire nkigen into the top-level uv workspace via pyproject.toml [tool.uv.workspace] members and [tool.uv.sources]. - Drop tests/debug, tests/outputs, tests/python (LIT) — internal scaffolding that doesn't belong in the open-source merge.
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
NKIgen is the NumPy-trace -> NISA-dialect lowering frontend (formerly "NKIPyKernelGen"). Imported as a top-level workspace member of the nkipy monorepo, replacing the
kernelgen/subfolder previously proposed in #55.Summary
This PR is a clean re-import addressing all of @vgene's feedback on #55. The diff goes from ~37K to ~32K lines (and from 230 to 176 files) by dropping internal scaffolding, deduplicating tests, and moving the Compiler Explorer wrapper out of scope.
Changes (vs. #55)
compiler_explorer/,.claude/skills/, internal scratch dirs (issues/,docs/,tools/kernel_agent/,docker/,examples/qwen3_embedding/, BrazilConfig). Compiler Explorer wrapper will land separately as a top-level component shared with other backends.tests/e2e/nkipy_tests/; the top-level/testsfolder is the single source of truth.Modeand tests; device is the golden, matching how nkipy's main backend tests are structured.nkipy_kernelgen→nkigen, harness decoratornkipy_kernelgen_test→nkigen_test.HardwareConstants.htarget-aware: replace globalMAX_PARTITION_DIM/MAX_FREE_DIM_MATMULwithgetSbufNumPartitions(target)/getMatmulFreeDimTileCap(target). The per-target values mirrornki.backends.mlir_tracer.target_infoin the publicnkiPython wheel.infer-layout,canonicalize-partition-dim, andlegalize-layouteach take a--target=option (defaulttrn2).pyproject.toml;setup.pyis now just the CMake build extension. Wirenkigeninto the top-level uv workspace viapyproject.toml[tool.uv.workspace] membersand[tool.uv.sources].tests/debug,tests/outputs,tests/python(LIT) — internal scaffolding that doesn't belong in the open-source merge.Test plan
pytest tests/passes/— 122 passedpytest tests/unit/— 128 passedpytest tests/e2e/— 52 passed, 10 skipped (HW-only on this trn1 box)Replaces
Closes #55.