Skip to content

[BACKEND] Fix CANN 8.5.1 bishengir compatibility #967

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Joiin0392 wants to merge 1 commit into
base: triton_v3.5.x
Choose a base branch
from
+25 −3
Open

[BACKEND] Fix CANN 8.5.1 bishengir compatibility #967

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions third_party/ascend/backend/compiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,18 @@ def make_ttir(mod, metadata, opt):
return mod


def _normalize_mlir_text_for_bishengir(linalg: str) -> str:
# CANN's bishengir tools cannot parse the newer ` to tensor<...>` result-type
# clause emitted by triton-opt for bufferization.to_tensor; strip it while
# keeping `restrict writable` (required by One-Shot Analysis).
if "bufferization.to_tensor" not in linalg:
return linalg
return "\n".join(
re.sub(r" to tensor<[^>]*>", "", line) if "bufferization.to_tensor" in line else line
for line in linalg.split("\n")
)


def ttir_to_linalg(mod, metadata, opt, *, named_ops=False):
# use triton_adapter to lower Triton-MLIR to linalg
# Get Triton-MLIR as string
Expand Down Expand Up @@ -169,7 +181,7 @@ def ttir_to_linalg(mod, metadata, opt, *, named_ops=False):
dump_manager = get_dump_manager(metadata["hash"])
dump_manager.put(str(mod), "kernel.ttadapter.mlir", binary=False)

return str(mod)
return _normalize_mlir_text_for_bishengir(str(mod))


def linalg_to_bc_by_triton_mlir_opt(linalg: str, metadata, opt):
Expand Down Expand Up @@ -712,7 +724,7 @@ def linalg_to_bin_enable_npu_compile_A2_A3(linalg: str, metadata, opt):
_compile_option_list += \
[f"--link-aicore-bitcode={bitcode}"]

_compile_option_list += [f"--link-aicore-bitcode={get_libdevice()}"]
pass

disable_size_align_for_cast = metadata["disable_size_align_for_cast"]
if disable_size_align_for_cast is not None:
Expand Down Expand Up @@ -866,7 +878,7 @@ class NPUOptions:
#
# If False, the compilation flow is:
# Linalg IR → LLIR → Binary (via bishengir-compile directly)
use_bytecode: bool = True
use_bytecode: bool = False
# take effect on the reorder instruction pattern for SIMT. The pattern is disabled by default.
enable_simt_reorder_instruction: bool = False
# disable simt fma optimization to get high precision
Expand Down
10 changes: 10 additions & 0 deletions third_party/ascend/language/cann/libdevice.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,16 @@ def ldexp(arg0, arg1, _semantic=None):

@core.extern
def pow(arg0, arg1, _semantic=None):
int_dtypes = (core.dtype("int1"), core.dtype("int8"), core.dtype("int16"),
core.dtype("int32"), core.dtype("int64"),
core.dtype("uint8"), core.dtype("uint16"),
core.dtype("uint32"), core.dtype("uint64"))
arg0 = _semantic.to_tensor(arg0)
arg1 = _semantic.to_tensor(arg1)
if arg0.dtype in int_dtypes:
arg0 = _semantic.cast(arg0, core.dtype("fp32"))
if arg1.dtype in int_dtypes:
arg1 = _semantic.cast(arg1, core.dtype("fp32"))
if triton_enable_libdevice_simt() and is_compile_on_910_95:
return core.extern_elementwise("", "", [arg0, arg1], {
(core.dtype("fp32"), core.dtype("fp32")): ("__hmf_pow_fp32", core.dtype("fp32")),
Expand Down