Skip to content

[AscendNPU-IR] Support the gather operator in Dev mode, and add test cases for the gather and not operators.#745

Open
lzp1021 wants to merge 6 commits intotile-ai:npuirfrom
lzp1021:supplement-api
Open

[AscendNPU-IR] Support the gather operator in Dev mode, and add test cases for the gather and not operators.#745
lzp1021 wants to merge 6 commits intotile-ai:npuirfrom
lzp1021:supplement-api

Conversation

@lzp1021
Copy link
Copy Markdown

@lzp1021 lzp1021 commented Apr 1, 2026

Support the gather operator in Dev mode, and add test cases for the gather and not operators.

@github-actions
Copy link
Copy Markdown

github-actions bot commented Apr 1, 2026

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

Comment thread src/target/codegen_npuir_dev.cc Outdated
Value dst = GenSubviewFromRegion(npuirop.dst, npuirop.dst_range);
Value indices = GenSubviewFromRegion(npuirop.indices, npuirop.indices_range);
Value src = GenExtractSliceFromRegion(npuirop.src, npuirop.src_range);
Value dst = GenExtractSliceFromRegion(npuirop.dst, npuirop.dst_range);
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not the right way. We don't use GenExtractSliceFromRegion for dst, but insert slice

Comment thread src/target/codegen_npuir_dev.cc Outdated
src, indices, dst);
auto gatherOp = builder.create<mlir::hivm::VGatherOp>(
builder.getUnknownLoc(), result_tensors, src, indices, dst);
SetVarValue(npuirop.dst, gatherOp->getResult(0));
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a if to find out if we need insert or not

@CeleNewYear CeleNewYear added the AscendNPU-IR Want to merge into the npuir branch label Apr 2, 2026
Comment on lines +38 to +39
A_VEC[:block_M, :block_N],
C_VEC[:block_M, :block_N],
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. now don't use npuir_gather
  2. This is not a full test 1) dynamic shape tile; 2) don't use all the shape (here :blockM = blockM)

C_VEC = T.alloc_shared([block_M, block_N], dtype)
T.copy(A[bx : bx + block_M, by : by + block_N], A_VEC)
T.copy(B[bx : bx + block_M, by : by + block_N], B_VEC)
T.gather(A_VEC, C_VEC, B_VEC)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will not test slice support of gather

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

AscendNPU-IR Want to merge into the npuir branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants